Skip to content

Dependency Update: Relax JAX Version Constraints#18

Merged
ZedongPeng merged 1 commit intoMIT-Lu-Lab:mainfrom
ZedongPeng:update_pyproject
May 8, 2025
Merged

Dependency Update: Relax JAX Version Constraints#18
ZedongPeng merged 1 commit intoMIT-Lu-Lab:mainfrom
ZedongPeng:update_pyproject

Conversation

@ZedongPeng
Copy link
Collaborator

@ZedongPeng ZedongPeng commented May 8, 2025

🔧 Dependency Update: Relax JAX Version Constraints

This PR updates the dependency versions in pyproject.toml to reflect recent improvements in the JAX ecosystem and ensure compatibility with macOS Intel hardware.

✅ Motivation

  • JAX PR #26876 fixes an issue where jax.numpy.linalg.norm would raise an error on empty arrays. (Fixed in JAX v0.5.3)

What happened in MPAX

In MPAX, this previously led to the following runtime error when the right-hand side array was empty:

  ValueError: zero-size array to reduction operation max which has no identity

Other considerations

Starting with jax==0.5.0, official wheels for macOS on Intel (x86) are no longer provided. To continue supporting users on macOS x86, we restrict the lower bound to >=0.4.38.

https://github.com/jax-ml/jax/releases/tag/jax-v0.5.0

@ZedongPeng ZedongPeng merged commit a823a63 into MIT-Lu-Lab:main May 8, 2025
10 checks passed
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