Skip to content

Fix native RHS preparation on JAX 0.10 - #180

Open
rogeriojorge wants to merge 3 commits into
en/custom_timestepperfrom
fix/native-rhs-jax-compat
Open

rogeriojorge wants to merge 3 commits into
en/custom_timestepperfrom
fix/native-rhs-jax-compat

Conversation

@rogeriojorge

@rogeriojorge rogeriojorge commented Oct 20, 2026 •

Copy link
Copy Markdown
Member

JAX 0.10 no longer re-exports partial_eval_jaxpr_nounits from jax.interpreters.partial_eval. Import the same implementation from its defining module; this remains an internal JAX dependency.

Replace the inherited max_steps source-string count with a behavioral check that Diffrax, native Dopri5, and native Dopri8 receive an explicit cap. The bounded-default assertion remains.

Validation: all three forwarding cases passed on JAX 0.11.2. Current-head Python 3.11/3.12 CI and patch coverage passed.

Based on en/custom_timestepper for integration into #169.

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.

1 participant