Add JAX dependency for Python 3.12 (#9438)

tengyifei · web-flow · commit 53ea5862afee · 2025-07-02T11:16:33.000-07:00
diff --git a/setup.py b/setup.py
@@ -403,7 +403,7 @@ def _get_jax_install_requirements():
   jax = f'jax @ https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/jax/jax-{_jax_version}-py3-none-any.whl'
 
   jaxlib = []
-  for python_minor_version in [9, 10, 11]:
+  for python_minor_version in [9, 10, 11, 12]:
     jaxlib.append(
         f'jaxlib @ https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/jaxlib/jaxlib-{_jaxlib_version}-cp3{python_minor_version}-cp3{python_minor_version}-manylinux2014_x86_64.whl ; python_version == "3.{python_minor_version}"'
     )

Original file line number	Diff line number	Diff line change
`@@ -403,7 +403,7 @@ def _get_jax_install_requirements():`
`403`	`403`	`jax = f'jax @ https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/jax/jax-{_jax_version}-py3-none-any.whl'`
`404`	`404`
`405`	`405`	`jaxlib = []`
`406`		`- for python_minor_version in [9, 10, 11]:`
	`406`	`+ for python_minor_version in [9, 10, 11, 12]:`
`407`	`407`	`jaxlib.append(`
`408`	`408`	`f'jaxlib @ https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/jaxlib/jaxlib-{_jaxlib_version}-cp3{python_minor_version}-cp3{python_minor_version}-manylinux2014_x86_64.whl ; python_version == "3.{python_minor_version}"'`
`409`	`409`	`)`