Bump jax to 0.7.0 #2131
Merged
Bump jax to 0.7.0 #2131
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…n of dependencies
rniczh
changed the base branch from
main
to
rniczh/bump-jax-to-0.7.0.dev20250703+cd1b9520b
JerryChen97
reviewed
Oct 23, 2025
….0.dev20250703+cd1b9520b
rniczh
changed the base branch from
rniczh/bump-jax-to-0.7.0.dev20250703+cd1b9520b
to
rniczh/update-llvm-version-20251009
…alyst into rniczh/bump-jax-to-0.7.0
Co-authored-by: Yushao Chen (Jerry) <chenys13@outlook.com>
Contributor
|
Fantastic, now even Codecov passes! @rniczh |
**Context:** Basically, for Catalyst, the new `bump-jax-patching` branch means there's no patching happenning by default anymore. In this way, we can better separate the impact of patching, instead of causing inheritance of patching from upstream to downstream. We could also better observe what patches are actually needed for other packages depending on Pennylane. **Description of the Change:** **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:**
Contributor
|
Now Catalyst's jax bumping branch should be independent of whether the PL Jax patches are global or local. |
JerryChen97
requested changes
Nov 21, 2025
Contributor
JerryChen97
left a comment
JerryChen97
left a comment
There was a problem hiding this comment.
Hold on temporarily until the PennyLane PR is also ready-to-merge
JerryChen97
approved these changes
Nov 28, 2025
Contributor
|
@rniczh Merge? |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Context:
Description of the Change:
lower_jaxpr_to_funhas changed, replacingpublicwithmain_function. Now, custom function names (e.g.,jit_…) were not permitted, withmainbeing the default. To address this,main_functionis now set to False, whilesym_visibilityis set to True for MLIR, retaining thejit_funcfunction name. Additionally,name_stackhas been removed. Do not propagate name_stacks into lower_jaxpr_to_fun. jax-ml/jax#29783make_varandget_varhas been removed. Should set the var (val) when initializing theDynamicJaxprTracer. Avoid strong refs to tracers in DynamicJaxprTrace. jax-ml/jax#29937. There is a typo in jax Fix typo in pjit.py jax-ml/jax#32718, we will patch the installation of jax in this PR.flatten_lowering_ir_argswithflatten_ir_values[JAX] Remove jax.interpreters.mlir.flatten_lowering_ir_args. jax-ml/jax#29706TracingEqnand replaceframe.eqns(originally usingJaxprEqn) withframe.tracing_eqns. roll-forward with fixes jax-ml/jax#30135pjit_phas been removed, usejit_pinsteadbin/patch_jax_installation.py)frontend/catalyst/jax_extras/patches.py)Benefits:
Possible Drawbacks:
Related GitHub Issues: