Merge pull request #20586 from superbobry:jaxlib

PiperOrigin-RevId: 624941598

Author: jax authors

CHANGELOG.md

@@ -29,7 +29,7 @@ Remember to align the itemized text with the first line of an item within a list
     is deprecated; empty inputs to softmax are now supported without setting this.
   * In {func}`jax.jit`, passing invalid `static_argnums` or `static_argnames`
     now leads to an error rather than a warning.
-
+  * The minimum jaxlib version is now 0.4.23.
 
 ## jaxlib 0.4.27
 
@@ -156,7 +156,7 @@ Remember to align the itemized text with the first line of an item within a list
       cannot interact, e.g., in arithmetic operations.
       Scopes are introduced by {func}`jax.experimental.jax2tf.convert`,
       {func}`jax.experimental.export.symbolic_shape`, {func}`jax.experimental.export.symbolic_args_specs`.
-      The scope of a symbolic expression `e` can be read with `e.scope` and passed 
+      The scope of a symbolic expression `e` can be read with `e.scope` and passed
       into the above functions to direct them to construct symbolic expressions in
       a given scope.
       See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#user-specified-symbolic-constraints.

jax/_src/dispatch.py

@@ -44,7 +44,6 @@
 from jax._src.interpreters import pxla
 from jax._src import lib
 from jax._src.lib import xla_client as xc
-from jax._src.lib import xla_extension_version
 from jax._src.monitoring import record_event_duration_secs
 from jax._src.partition_spec import PartitionSpec
 from jax._src.sharding import Sharding
@@ -82,15 +81,11 @@ def apply_primitive(prim, *args, **params):
   fun = xla_primitive_callable(prim, **params)
   # TODO(yashkatariya): Investigate adding is_primitive to jit and never
   # triggering the disable jit path instead of messing around with it here.
-  if xla_extension_version >= 218:
-    prev = lib.jax_jit.swap_thread_local_state_disable_jit(False)
-    try:
-      outs = fun(*args)
-    finally:
-      lib.jax_jit.swap_thread_local_state_disable_jit(prev)
-  else:
-    with config.disable_jit(False):
-      outs = fun(*args)
+  prev = lib.jax_jit.swap_thread_local_state_disable_jit(False)
+  try:
+    outs = fun(*args)
+  finally:
+    lib.jax_jit.swap_thread_local_state_disable_jit(prev)
   return outs
 
 @util.cache()

jax/_src/interpreters/mlir.py

@@ -49,7 +49,6 @@
 from jax._src.layout import AutoLayout, DeviceLocalLayout
 from jax._src.lib import xla_client as xc
 from jax._src.lib import xla_extension
-from jax._src.lib import xla_extension_version
 from jax._src.lib.mlir import dialects
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import func as func_dialect
@@ -911,7 +910,7 @@ def lower_jaxpr_to_module(
         "In multi-platform lowering either all or no lowering platforms "
         f"should support donation. Lowering for {platforms} of which "
         f"only {platforms_with_donation} support donation")
-    if num_partitions > 1 and xla_extension_version >= 220 and (
+    if num_partitions > 1 and (
         result_shardings is None or all(s is None for s in result_shardings)):
       xla_donated_args = donated_args
     if xla_donated_args is None:

jax/_src/interpreters/pxla.py

@@ -2918,12 +2918,8 @@ def from_hlo(name: str,
         in_shardings, out_shardings, committed, da = _get_metadata_jit_pmap(
             xla_executable.local_devices(), len(in_shardings), len(out_shardings))
 
-    if xla_extension_version >= 217:
-      in_layouts, out_layouts = _get_layouts_from_executable(
-          xla_executable, in_layouts, out_layouts, len(ordered_effects))
-    else:
-      assert all(i is None for i in in_layouts)
-      assert all(o is None for o in out_layouts)
+    in_layouts, out_layouts = _get_layouts_from_executable(
+        xla_executable, in_layouts, out_layouts, len(ordered_effects))
 
     out_shardings = maybe_recover_user_shardings(
         in_shardings, out_shardings, global_in_avals, global_out_avals)

jax/_src/lax/fft.py

@@ -30,8 +30,6 @@
 from jax._src.interpreters import mlir
 from jax._src.lib.mlir.dialects import hlo
 from jax._src.lib import xla_client
-from jax._src.lib import xla_extension_version
-from jax._src.lib import ducc_fft
 from jax._src.numpy.util import promote_dtypes_complex, promote_dtypes_inexact
 
 __all__ = [
@@ -122,76 +120,6 @@ def _fft_lowering(ctx, x, *, fft_type, fft_lengths):
   ]
 
 
-def _fft_lowering_cpu(ctx, x, *, fft_type, fft_lengths):
-  x_aval, = ctx.avals_in
-
-  in_shape = x_aval.shape
-  dtype = x_aval.dtype
-  out_aval, = ctx.avals_out
-  out_shape = out_aval.shape
-
-  forward = fft_type in (xla_client.FftType.FFT, xla_client.FftType.RFFT)
-  ndims = len(in_shape)
-  assert len(fft_lengths) >= 1
-  assert len(fft_lengths) <= ndims, (fft_lengths, ndims)
-  assert len(in_shape) == len(out_shape) == ndims
-
-  # PocketFft does not allow size 0 dimensions.
-  if 0 in in_shape or 0 in out_shape:
-    if fft_type == xla_client.FftType.RFFT:
-      assert dtype in (np.float32, np.float64), dtype
-      out_dtype = np.dtype(np.complex64 if dtype == np.float32 else np.complex128)
-
-    elif fft_type == xla_client.FftType.IRFFT:
-      assert np.issubdtype(dtype, np.complexfloating), dtype
-      out_dtype = np.dtype(np.float32 if dtype == np.complex64 else np.float64)
-
-    else:
-      assert np.issubdtype(dtype, np.complexfloating), dtype
-      out_dtype = dtype
-
-    zero = mlir.ir_constant(np.array(0, dtype=out_dtype))
-    return [
-        mlir.broadcast_in_dim(ctx, zero, out_aval, broadcast_dimensions=[])]
-
-  strides_in = []
-  stride = 1
-  for d in reversed(in_shape):
-    strides_in.append(stride)
-    stride *= d
-  strides_in = mlir.shape_tensor(
-      mlir.eval_dynamic_shape(ctx, tuple(reversed(strides_in))))
-
-  strides_out = []
-  stride = 1
-  for d in reversed(out_shape):
-    strides_out.append(stride)
-    stride *= d
-  strides_out = mlir.shape_tensor(
-      mlir.eval_dynamic_shape(ctx, tuple(reversed(strides_out))))
-
-  # scale = 1. if forward else (1. / np.prod(fft_lengths)) as a f64[1] tensor
-  double_type = mlir.ir.RankedTensorType.get((), mlir.ir.F64Type.get())
-  size_fft_length_prod = np.prod(fft_lengths) if fft_lengths else 1
-  size_fft_lengths, = mlir.eval_dynamic_shape_as_vals(ctx, (size_fft_length_prod,))
-  size_fft_lengths = hlo.ConvertOp(double_type, size_fft_lengths)
-  one = mlir.ir_constant(np.float64(1.))
-  scale = one if forward else hlo.DivOp(one, size_fft_lengths)
-  scale = hlo.ReshapeOp(
-      mlir.ir.RankedTensorType.get((1,), mlir.ir.F64Type.get()),
-      scale).result
-
-  in_shape = mlir.shape_tensor(mlir.eval_dynamic_shape(ctx, in_shape))
-  out_shape = mlir.shape_tensor(mlir.eval_dynamic_shape(ctx, out_shape))
-  in_shape = in_shape if fft_type != xla_client.FftType.IRFFT else out_shape
-
-  result_type = mlir.aval_to_ir_type(out_aval)
-  return [ducc_fft.dynamic_ducc_fft_hlo(
-      result_type, x,
-      input_dtype=x_aval.dtype, ndims=ndims, input_shape=in_shape,
-      strides_in=strides_in, strides_out=strides_out, scale=scale,
-      fft_type=fft_type, fft_lengths=fft_lengths, result_shape=out_shape)]
-
 def _naive_rfft(x, fft_lengths):
   y = fft(x, xla_client.FftType.FFT, fft_lengths)
   n = fft_lengths[-1]
@@ -253,8 +181,3 @@ def _fft_batching_rule(batched_args, batch_dims, fft_type, fft_lengths):
 mlir.register_lowering(fft_p, _fft_lowering)
 ad.deflinear2(fft_p, _fft_transpose_rule)
 batching.primitive_batchers[fft_p] = _fft_batching_rule
-
-# TODO(phawkins): when jaxlib 0.4.21 is the minimum, use XLA's FFT lowering
-# always on CPU. At that point, we can also delete the DUCC FFT kernel from JAX.
-if xla_extension_version < 211:
-  mlir.register_lowering(fft_p, _fft_lowering_cpu, platform='cpu')

jax/_src/xla_bridge.py

@@ -129,10 +129,6 @@ def _get_tpu_library_path() -> str | None:
 
   libtpu_module = maybe_import_libtpu()
   if libtpu_module is not None:
-    if xla_extension_version < 212:
-      # xla_extension_version < 212 uses tpu_tracer which requires calling
-      # configure_library_path.
-      libtpu_module.configure_library_path()
     return libtpu_module.get_library_path()
 
   return None
@@ -226,28 +222,17 @@ def register_backend_factory(name: str, factory: BackendFactory, *,
 
 
 def make_cpu_client() -> xla_client.Client:
-  if xla_extension_version >= 223:
-    collectives: xla_client._xla.CpuCollectives | None = None
-    if _CPU_ENABLE_GLOO_COLLECTIVES.value:
-      collectives = xla_client._xla.make_gloo_tcp_collectives(  # type: ignore
-        distributed_client=distributed.global_state.client,
-      )
-    return xla_client.make_cpu_client(  # type: ignore
+  collectives: xla_client._xla.CpuCollectives | None = None
+  if _CPU_ENABLE_GLOO_COLLECTIVES.value:
+    collectives = xla_client._xla.make_gloo_tcp_collectives(  # type: ignore
       distributed_client=distributed.global_state.client,
-      node_id=distributed.global_state.process_id,
-      num_nodes=distributed.global_state.num_processes,
-      collectives=collectives,
     )
-  elif xla_extension_version >= 216:
-    # TODO(phawkins): remove type: ignore after updating jaxlib version used for
-    # mypy checks.
-    return xla_client.make_cpu_client(  # type: ignore
-      distributed_client=distributed.global_state.client,
-      node_id=distributed.global_state.process_id,
-      num_nodes=distributed.global_state.num_processes,
-    )
-  else:
-    return xla_client.make_cpu_client()
+  return xla_client.make_cpu_client(  # type: ignore
+    distributed_client=distributed.global_state.client,
+    node_id=distributed.global_state.process_id,
+    num_nodes=distributed.global_state.num_processes,
+    collectives=collectives,
+  )
 
 
 register_backend_factory(

jax/version.py

@@ -133,7 +133,7 @@ def make_release_tree(self, base_dir, files):
 
 
 __version__ = _get_version_string()
-_minimum_jaxlib_version = "0.4.20"
+_minimum_jaxlib_version = "0.4.23"
 
 def _version_as_tuple(version_str):
   return tuple(int(i) for i in version_str.split(".") if i.isdigit())

tests/api_test.py

@@ -58,7 +58,6 @@
 from jax._src.interpreters import partial_eval as pe
 from jax._src.lib import xla_client
 from jax._src.lib import xla_extension
-from jax._src.lib import xla_extension_version
 import jax._src.util as jax_util
 from jax.ad_checkpoint import checkpoint_name, checkpoint as new_checkpoint
 import jax.custom_batching
@@ -4590,7 +4589,6 @@ def test_mesh_creation_error_message(self):
     with self.assertRaisesRegex(ValueError, "ndim of its first argument"):
       jax.sharding.Mesh(jax.devices(), ("x", "y"))
 
-  @unittest.skipIf(xla_extension_version < 222, 'jaxlib version too old')
   def test_jit_boundmethod_reference_cycle(self):
     class A:
       def __init__(self):

tests/array_interoperability_test.py

@@ -221,7 +221,6 @@ def testJaxToNumpy(self, shape, dtype):
     x_np = np.from_dlpack(x_jax)
     self.assertAllClose(x_np, x_jax)
 
-  @unittest.skipIf(xla_extension_version < 221, "Requires newer jaxlib")
   def testNondefaultLayout(self):
     # Generate numpy array with nonstandard layout
     a = np.arange(4).reshape(2, 2)

tests/export_back_compat_test.py

@@ -60,7 +60,6 @@
 
 from jax._src import config
 from jax._src import test_util as jtu
-from jax._src.lib import version as jaxlib_version
 
 config.parse_flags_with_absl()
 
@@ -142,23 +141,16 @@ def test_ducc_fft(self):
     def func(x):
       return lax.fft(x, fft_type="fft", fft_lengths=(4,))
 
-    # An old lowering, with ducc_fft. We keep it for 6 months.
-    data = self.load_testdata(cpu_ducc_fft.data_2023_03_17)
-    if jaxlib_version <= (0, 4, 20):
-      expect_current_custom_calls = ["dynamic_ducc_fft"]
-    else:
-      # We have changed the lowering for fft since we saved this data.
-      # FFT no longer lowers to a custom call.
-      expect_current_custom_calls = []
+    # TODO(b/311175955): Remove this test and the corresponding custom calls.
 
-    self.run_one_test(func, data,
-                      expect_current_custom_calls=expect_current_custom_calls)
+    # An old lowering, with ducc_fft.
+    data = self.load_testdata(cpu_ducc_fft.data_2023_03_17)
+    self.run_one_test(func, data, expect_current_custom_calls=[])
 
     # A newer lowering, with dynamic_ducc_fft.
     data = self.load_testdata(cpu_ducc_fft.data_2023_06_14)
     # FFT no longer lowers to a custom call.
-    self.run_one_test(func, data,
-                      expect_current_custom_calls=expect_current_custom_calls)
+    self.run_one_test(func, data, expect_current_custom_calls=[])
 
   def cholesky_input(self, shape, dtype):
     a = jtu.rand_default(self.rng())(shape, dtype)