[JAX] Update JAX version requirement in pyproject.toml (#2197)

update jax requirements Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

[JAX] Update JAX version requirement in pyproject.toml (#2197)
update jax requirements Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
afd15a16 · Phuong Nguyen · GitHub · 3f875fb5 · afd15a16
Unverified Commit afd15a16 authored Sep 24, 2025 by Phuong Nguyen Committed by GitHub Sep 24, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

pyproject.toml pyproject.toml +2 -1

No files found.
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,8 @@
 # See LICENSE for license information.

 [build-system]
-requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax[cuda12]", "flax>=0.7.1"]
+requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip",
+"torch>=2.1", "jax>=0.5.0", "flax>=0.7.1"]

 # Use legacy backend to import local packages in setup.py
 build-backend = "setuptools.build_meta:__legacy__"