Skip to content

Require JAX 0.6+ for import compatibility#56

Open
talolard wants to merge 1 commit into
brentyi:mainfrom
talolard:tal/fix-jax-extend-core
Open

Require JAX 0.6+ for import compatibility#56
talolard wants to merge 1 commit into
brentyi:mainfrom
talolard:tal/fix-jax-extend-core

Conversation

@talolard
Copy link
Copy Markdown

@talolard talolard commented May 2, 2026

Summary

Fixes #55 for the current main package line by aligning package metadata with the JAX API era that no longer exposes internals through jax.core.

Current main no longer contains the legacy jaxls/typed.py code path that referenced jax.core.JaxprEqn; a fresh import succeeds with JAX 0.6+. This PR makes that compatibility explicit by:

  • raising the runtime dependency floor to jax>=0.6.0 and jaxlib>=0.6.0
  • adding a subprocess import regression test so import-time failures from removed JAX internals are caught directly

Validation

  • make format fails before running because Makefile references an undefined examples extra: error: Extra examples is not defined in the project's optional-dependencies table
  • uvx ruff format .
  • uvx ruff check .
  • uv run --python 3.12 --with-editable . --with 'jax==0.6.0' --with 'jaxlib==0.6.0' python - <<'PY' ... import jax; import jaxls ... PY
  • uv run --python 3.12 --extra dev --extra cholmod pytest
  • uv run --python 3.12 --extra dev --extra docs --extra cholmod pyright .
  • uv run --python 3.12 --extra dev python transpile_py310.py && git diff --quiet src/jaxls/_py310/
  • uvx ruff check . && uvx ruff format --check .
  • for py in 3.10 3.11 3.12 3.13; do uv run --python "$py" --extra dev --extra cholmod pytest || exit $?; done

Note: local cholmod validation required installing the same system dependency used by CI: libsuitesparse-dev.

Current main no longer uses the legacy typed.py path that referenced jax.core.JaxprEqn, but the package metadata still allowed older pre-0.6 JAX installs. Raise the JAX and jaxlib floors to the version where moved internals live under jax.extend.core and add a subprocess import regression test for issue brentyi#55.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Import fails with JAX >= 0.6.0: jax.core.JaxprEqn removed

2 participants