Only require jax to build where jaxlib has a wheel

[jdeast]

The JAX extension is already optional in CMakeLists.txt -- it probes for jax.ffi.include_dir() and, when that fails, prints Skipping JAX extension and builds everything else. But pip resolves [build-system] requires in an isolated environment before CMake runs, so that graceful path is unreachable and the pinned jax==0.8.0 makes the extension mandatory in practice.

Where jaxlib publishes no wheel, that makes the sdist unbuildable. macOS x86_64 is the case that bites: jaxlib's last wheel there is 0.4.38 and jaxlib ships no sdist, so the build fails while installing build dependencies, before any C++ is compiled:

× installing build dependencies for celerite2 did not run successfully.
  Collecting jax==0.8.0
  ERROR: Could not find a version that satisfies the requirement jaxlib<=0.8.0,>=0.8.0
         (from versions: 0.4.18, ..., 0.4.36, 0.4.38)
  ERROR: No matching distribution found for jaxlib<=0.8.0,>=0.8.0

This is not fixable downstream either: PIP_CONSTRAINT only narrows a requirement, so it cannot relax an == pin.

Why a marker rather than just dropping jax

Dropping jax from requires would work for the sdist but would silently change your releases. wheels.yml gets jax only via [build-system] requires, so without it CMake would take the "Skipping JAX extension" branch during cibuildwheel and the published wheels would quietly stop containing celerite2/jax/xla_ops -- building, installing, and passing CI exactly as before.

The marker keeps jax present on every platform you build wheels for and drops it only where jaxlib cannot be installed at all.

Verification

Run on a fork with this change:

• Wheels: all three cibuildwheel jobs green, and all 15 wheels (macOS arm64, manylinux x86_64, win_amd64; cp311-cp314) contain celerite2/jax/xla_ops -- checked by inspecting the artifacts, since a missing extension would not otherwise show up as a failure.

• Python: all 5 sessions green (core, jax, pymc, pymc_jax, tests-pymc).

• Marker evaluated directly against the wheel matrix:

  platform	wheels built?	jax required
  linux x86_64	yes	True
  linux aarch64	yes	True
  macOS arm64	yes	True
  windows AMD64	yes	True
  macOS x86_64	no	False

Found while getting a downstream package installable on Intel macOS. Happy to adjust the marker if you would rather express it differently.

🤖 Generated with Claude Code