Skip to content

Conversation

@cmj2002
Copy link
Collaborator

@cmj2002 cmj2002 commented Oct 18, 2025

This PR addresses a dependency conflict that can occur when users have jax==0.5.3 installed and do not specify an upper bound for the orbax-checkpoint dependency.

The issue is as follows:

  1. flow-rl currently requires jax==0.5.3 and does not specify an upper bound for the orbax-checkpoint dependency.
  2. orbax-checkpoint version 0.11.24 introduced a call to jax.monitoring.record_scalar, which only exists in jax>=0.6.0.
  3. orbax-checkpoint version 0.11.24 did not enforce a jax>=0.6.0 dependency, allowing it to be installed alongside jax==0.5.3.
  4. Although orbax-checkpoint version 0.11.25 corrected this by adding the jax>=0.6.0 requirement, version 0.11.24 remains available on PyPI, leading to installation and runtime failures for users with the current dependency setup of flow-rl.

Solution: This PR pins the orbax-checkpoint dependency to 0.11.23, the newest version which does not require jax>=0.6.0.

@typoverflow
Copy link
Owner

typoverflow commented Oct 18, 2025

Is it possible for us to upgrade our jax version to >= 0.6.0?

@typoverflow typoverflow merged commit 414d090 into master Oct 19, 2025
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.

3 participants