Repository: jax-ml/jax
Pull request #27502 (2025-03-26 22:39:54)

Rename xla_extension_version to jaxlib_extension_version to reflect its new
Rename xla_extension_version to jaxlib_extension_version to reflect its new
scope.



Patch:
diff --git a/docs/jep/9419-jax-versioning.md b/docs/jep/9419-jax-versioning.md
index b964aa2af..85b95257e 100644
--- a/docs/jep/9419-jax-versioning.md
+++ b/docs/jep/9419-jax-versioning.md
@@ -167,16 +167,16 @@ We maintain an additional version number (`_version`) in
 [`xla_client.py` in the XLA repository](https://github.com/openxla/xla/blob/main/xla/python/xla_client.py).
 The idea is that this version number, is defined in `xla/python`
 together with the C++ parts of JAX, is also accessible to JAX Python as
-`jax._src.lib.xla_extension_version`, and must
+`jax._src.lib.jaxlib_extension_version`, and must
 be incremented every time that a change is made to the XLA/Python code that has
 backwards compatibility implications for `jax`. The JAX Python code can then use
 this version number to maintain backwards compatibility, e.g.:
 
 ```
-from jax._src.lib import xla_extension_version
+from jax._src.lib import jaxlib_extension_version
 
 # 123 is the new version number for _version in xla_client.py
-if xla_extension_version >= 123:
+if jaxlib_extension_version >= 123:
   # Use new code path
   ...
 else:
diff --git a/jax/_src/callback.py b/jax/_src/callback.py
index dc60bfb94..683da6663 100644
--- a/jax/_src/callback.py
+++ b/jax/_src/callback.py
@@ -33,7 +33,7 @@ from jax._src import sharding_impls
 from jax._src import tree_util
 from jax._src import util
 from jax._src import xla_bridge as xb
-from jax._src.lib import xla_extension_version
+from jax._src.lib import jaxlib_extension_version
 from jax._src.interpreters import ad
 from jax._src.interpreters import batching
 from jax._src.interpreters import mlir
@@ -827,7 +827,7 @@ def emit_python_callback(
     return outputs, token, None
 
   # TODO(dsuo): Remove this once we bump minimum_jaxlib_version to "0.5.4".
-  if xla_extension_version <= 320:
+  if jaxlib_extension_version <= 320:
     result_types = mlir.flatten_ir_types([mlir.aval_to_ir_type(aval) for aval in result_avals])
     if token:
 
diff --git a/jax/_src/lib/__init__.py b/jax/_src/lib/__init__.py
index be551449a..fef5d2c26 100644
--- a/jax/_src/lib/__init__.py
+++ b/jax/_src/lib/__init__.py
@@ -117,7 +117,7 @@ import jaxlib.gpu_linalg as gpu_linalg  # pytype: disable=import-error  # noqa:
 # Only for the internal usage of the JAX developers, we expose a version
 # number that can be used to perform changes without breaking the main
 # branch on the Jax github.
-xla_extension_version: int = getattr(xla_client, '_version', 0)
+jaxlib_extension_version: int = getattr(xla_client, '_version', 0)
 
 import jaxlib.gpu_rnn as gpu_rnn  # pytype: disable=import-error  # noqa: F401
 import jaxlib.gpu_triton as gpu_triton # pytype: disable=import-error  # noqa: F401
diff --git a/jaxlib/xla/xla_client.py b/jaxlib/xla/xla_client.py
index eb7a3b575..80cdeef47 100644
--- a/jaxlib/xla/xla_client.py
+++ b/jaxlib/xla/xla_client.py
@@ -49,7 +49,7 @@ ops = _xla.ops
 profiler = _xla.profiler
 
 # Just an internal arbitrary increasing number to help with backward-compatible
-# changes. In JAX, reference this via jax._src.lib.xla_extension_version.
+# changes. In JAX, reference this via jax._src.lib.jaxlib_extension_version.
 _version = 322
 
 # Version number for MLIR:Python components.


Test patch:
diff --git a/tests/python_callback_test.py b/tests/python_callback_test.py
index 5650a2d4f..a8442b4a1 100644
--- a/tests/python_callback_test.py
+++ b/tests/python_callback_test.py
@@ -28,7 +28,7 @@ from jax._src import core
 from jax._src import dispatch
 from jax._src import test_util as jtu
 from jax._src import util
-from jax._src.lib import xla_extension_version
+from jax._src.lib import jaxlib_extension_version
 from jax.experimental import io_callback
 from jax.experimental import pjit
 from jax.experimental.shard_map import shard_map
@@ -588,8 +588,8 @@ class PythonCallbackTest(jtu.JaxTestCase):
 
   @parameterized.parameters("int2", "int4", "uint2", "uint4")
   def test_subbyte_operands(self, dtype: str):
-    if xla_extension_version <= 321:
-      self.skipTest("Requires xla_extension_version >= 322.")
+    if jaxlib_extension_version <= 321:
+      self.skipTest("Requires jaxlib_extension_version >= 322.")
     def get(x):
       return x
     def f(x):
@@ -613,8 +613,8 @@ class PythonCallbackTest(jtu.JaxTestCase):
 
   @parameterized.parameters("int2", "int4", "uint2", "uint4")
   def test_subbyte_results(self, dtype: str):
-    if xla_extension_version <= 321:
-      self.skipTest("Requires xla_extension_version >= 322.")
+    if jaxlib_extension_version <= 321:
+      self.skipTest("Requires jaxlib_extension_version >= 322.")
     def get():
       return np.arange(8, dtype=dtype)