Remove code that existed to support jaxlib < 0.4.32.

New minimum versions:
* jaxlib 0.4.32
* xla_extension_version 283
* mlir_api_version 57

PiperOrigin-RevId: 673526233

jax/_src/compiler.py

@@ -33,7 +33,6 @@
 from jax._src import traceback_util
 from jax._src.interpreters import mlir
 from jax._src.lib import xla_client as xc
-from jax._src.lib import xla_extension_version
 from jax._src.lib.mlir import ir
 import numpy as np
 
@@ -157,8 +156,7 @@ def get_compile_options(
  build_options = compile_options.executable_build_options
  build_options.use_spmd_partitioning = use_spmd_partitioning
  build_options.use_auto_spmd_partitioning = use_auto_spmd_partitioning
- if xla_extension_version >= 280:
- build_options.use_shardy_partitioner = use_shardy_partitioner
+ build_options.use_shardy_partitioner = use_shardy_partitioner
  if fdo_profile is not None:
  build_options.fdo_profile = fdo_profile
  if use_auto_spmd_partitioning:

jax/_src/interpreters/pxla.py

@@ -61,7 +61,6 @@
 from jax._src.interpreters import xla
 from jax._src.layout import DeviceLocalLayout, AutoLayout, Layout
 from jax._src.lib import xla_client as xc
-from jax._src.lib import xla_extension_version
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo
 from jax._src.partition_spec import PartitionSpec
@@ -3022,12 +3021,8 @@ def aot_cache_miss(*args, **kwargs):
  self.unsafe_call.name, None, aot_cache_miss, [], [], [],
  tree_util.dispatch_registry, cc_shard_arg)
 
-if xla_extension_version < 282:
- def cc_shard_arg(x, sharding):
- return shard_args([sharding], [None], [x])[0]
-else:
- def cc_shard_arg(x, sharding, layout): # type: ignore
- return shard_args([sharding], [layout], [x])[0]
+def cc_shard_arg(x, sharding, layout): # type: ignore
+ return shard_args([sharding], [layout], [x])[0]
 
 
 def check_arg_avals_for_call(ref_avals, arg_avals,

jax/_src/lax/linalg.py

@@ -514,11 +514,7 @@ def _cholesky_cpu_lowering(ctx, operand):
  out_aval, = ctx.avals_out
  batch_dims = operand_aval.shape[:-2]
  op_shape_vals = mlir.eval_dynamic_shape_as_ivals(ctx, operand_aval.shape)
- # TODO(b/344892332): Remove the check after the compatibility period.
- if jaxlib_version < (0, 4, 31):
- ctx_arg = ()
- else:
- ctx_arg = (ctx,)
+ ctx_arg = (ctx,)
  result, info = lapack.potrf_hlo(*ctx_arg, operand_aval.dtype, operand,
  lower=True, a_shape_vals=op_shape_vals)
 
@@ -556,7 +552,7 @@ def _cholesky_update_abstract_eval(r_matrix, w_vector):
 
 def _cholesky_update_gpu_lowering_rule(target_name_prefix, ctx, r_matrix, w_vector):
  # TODO(b/360781533): Remove guard after 3 week forward compatibility period.
- if ctx.is_forward_compat() or jaxlib_version < (0, 4, 32):
+ if ctx.is_forward_compat():
  r_matrix_aval, _ = ctx.avals_in
  try:
  [platform] = ctx.module_context.platforms
@@ -726,8 +722,7 @@ def _eig_cpu_lowering(ctx, operand, *, compute_left_eigenvectors,
  out_aval = ctx.avals_out[0]
  batch_dims = operand_aval.shape[:-2]
  op_shape_vals = mlir.eval_dynamic_shape_as_ivals(ctx, operand_aval.shape)
- # TODO(b/344892332): Remove the conditional after the compatibility period.
- ctx_args = (ctx,) if jaxlib_version >= (0, 4, 32) else ()
+ ctx_args = (ctx,)
  w, vl, vr, info = lapack.geev_hlo(*ctx_args, operand_aval.dtype, operand,
  input_shape_vals=op_shape_vals,
  jobvl=compute_left_eigenvectors,
@@ -937,8 +932,7 @@ def _eigh_cpu_gpu_lowering(
  op_shape_vals = mlir.eval_dynamic_shape_as_ivals(ctx, operand_aval.shape)
  cpu_args = []
  if platform == "cpu":
- # TODO(b/344892332): Remove the conditional after the compatibility period.
- ctx_args = (ctx,) if jaxlib_version >= (0, 4, 32) else ()
+ ctx_args = (ctx,)
  cpu_args.extend(ctx_args)
  v, w, info = syevd_impl(*cpu_args, operand_aval.dtype, operand,
  a_shape_vals=op_shape_vals, lower=lower)
@@ -1511,9 +1505,9 @@ def _lu_cpu_gpu_lowering(getrf_impl, ctx, operand, *, platform: str,
  info_aval = ShapedArray(batch_dims, np.dtype(np.int32))
  m = operand_aval.shape[-2]
 
- # TODO(b/357034884): Remove version gate once jaxlib 0.4.32 is the minimum
- # version and the forward compat flag after the 3 week compatibility window.
- if jaxlib_version < (0, 4, 32) or ctx.is_forward_compat():
+ # TODO(b/357034884): Remove version gate on the forward compat flag after the
+ # 3 week compatibility window.
+ if ctx.is_forward_compat():
  if not is_constant_shape(operand_aval.shape[-2:]):
  raise NotImplementedError(
  "Shape polymorphism for native lowering for lu on CPU and GPU is "
@@ -1757,9 +1751,8 @@ def _geqrf_cpu_gpu_lowering(geqrf_impl, batched_geqrf_impl, ctx, a, *,
  a_out, taus, info_geqrf = geqrf_impl(a_aval.dtype, a)
  else:
  a_shape_vals = mlir.eval_dynamic_shape_as_ivals(ctx, a_aval.shape)
- # TODO(b/344892332): Remove the conditional after the compatibility period
  ctx_args = (
- (ctx,) if platform == "cpu" and jaxlib_version >= (0, 4, 32) else ()
+ (ctx,) if platform == "cpu" else ()
  )
  a_out, taus, *maybe_info_geqrf = geqrf_impl(
  *ctx_args, a_aval.dtype, a, a_shape_vals=a_shape_vals
@@ -1881,9 +1874,8 @@ def _householder_product_cpu_gpu_lowering(orgqr_impl, ctx, a, taus, *,
  f"on GPU is not implemented; b/261671778; {a_aval.shape}")
  a, info_orgqr = orgqr_impl(a_aval.dtype, a, taus)
  else:
- # TODO(b/344892332): Remove the conditional after the compatibility period
  ctx_args = (
- (ctx,) if platform == "cpu" and jaxlib_version >= (0, 4, 32) else ()
+ (ctx,) if platform == "cpu" else ()
  )
  a_shape_vals = mlir.eval_dynamic_shape_as_ivals(ctx, a_aval.shape)
  tau_shape_vals = mlir.eval_dynamic_shape_as_ivals(ctx, taus_aval.shape)
@@ -2152,8 +2144,7 @@ def _svd_cpu_gpu_lowering(
  compute_uv=compute_uv)
  else:
  a_shape_vals = mlir.eval_dynamic_shape_as_ivals(ctx, operand_aval.shape)
- # TODO(b/344892332): Remove the conditional after the compatibility period.
- ctx_args = (ctx,) if jaxlib_version >= (0, 4, 32) else ()
+ ctx_args = (ctx,)
  s, u, vt, info = gesvd_impl(*ctx_args, operand_aval.dtype, operand,
  full_matrices=full_matrices,
  compute_uv=compute_uv,

tests/export_back_compat_test.py

@@ -66,7 +66,6 @@
 from jax._src import config
 from jax._src import test_util as jtu
 from jax._src.lib import cuda_versions
-from jax._src.lib import version as jaxlib_version
 
 config.parse_flags_with_absl()
 
@@ -190,14 +189,11 @@ def test_cpu_cholesky_lapack_potrf(self, dtype_name="f32"):
  atol = dict(f32=1e-4, f64=1e-12, c64=1e-4, c128=1e-12)[dtype_name]
 
  data = self.load_testdata(cpu_cholesky_lapack_potrf.data_2023_06_19[dtype_name])
- # TODO(b/344892332): Remove the check after the compatibility period.
- has_xla_ffi_support = jaxlib_version >= (0, 4, 31)
  self.run_one_test(func, data, rtol=rtol, atol=atol)
- if has_xla_ffi_support:
- with config.export_ignore_forward_compatibility(True):
- # FFI Kernel test
- data = self.load_testdata(cpu_cholesky_lapack_potrf.data_2024_05_31[dtype_name])
- self.run_one_test(func, data, rtol=rtol, atol=atol)
+ with config.export_ignore_forward_compatibility(True):
+ # FFI Kernel test
+ data = self.load_testdata(cpu_cholesky_lapack_potrf.data_2024_05_31[dtype_name])
+ self.run_one_test(func, data, rtol=rtol, atol=atol)
 
  @parameterized.named_parameters(
  dict(testcase_name=f"_dtype={dtype_name}", dtype_name=dtype_name)
@@ -258,14 +254,11 @@ def check_eigenvalue_is_in_array(eigenvalue, eigenvalues_array):
 
  self.run_one_test(func, data, rtol=rtol, atol=atol,
  check_results=check_eig_results)
- # TODO(b/344892332): Remove the check after the compatibility period.
- has_xla_ffi_support = jaxlib_version >= (0, 4, 32)
- if has_xla_ffi_support:
- with config.export_ignore_forward_compatibility(True):
- # FFI Kernel test
- data = self.load_testdata(cpu_eig_lapack_geev.data_2024_08_19[dtype_name])
- self.run_one_test(func, data, rtol=rtol, atol=atol,
- check_results=check_eig_results)
+ with config.export_ignore_forward_compatibility(True):
+ # FFI Kernel test
+ data = self.load_testdata(cpu_eig_lapack_geev.data_2024_08_19[dtype_name])
+ self.run_one_test(func, data, rtol=rtol, atol=atol,
+ check_results=check_eig_results)
 
  @staticmethod
  def eigh_input(shape, dtype):
@@ -316,14 +309,11 @@ def test_cpu_eigh_lapack_syevd(self, dtype_name="f32"):
  atol = dict(f32=1e-4, f64=1e-12, c64=1e-4, c128=1e-12)[dtype_name]
  self.run_one_test(func, data, rtol=rtol, atol=atol,
  check_results=partial(self.check_eigh_results, operand))
- # TODO(b/344892332): Remove the check after the compatibility period.
- has_xla_ffi_support = jaxlib_version >= (0, 4, 32)
- if has_xla_ffi_support:
- # FFI Kernel test
- with config.export_ignore_forward_compatibility(True):
- data = self.load_testdata(cpu_eigh_lapack_syev.data_2024_08_19[dtype_name])
- self.run_one_test(func, data, rtol=rtol, atol=atol,
- check_results=partial(self.check_eigh_results, operand))
+ # FFI Kernel test
+ with config.export_ignore_forward_compatibility(True):
+ data = self.load_testdata(cpu_eigh_lapack_syev.data_2024_08_19[dtype_name])
+ self.run_one_test(func, data, rtol=rtol, atol=atol,
+ check_results=partial(self.check_eigh_results, operand))
 
  @parameterized.named_parameters(
  dict(testcase_name=f"_dtype={dtype_name}_{variant}",
@@ -385,8 +375,6 @@ def test_cuda_lu_pivots_to_permutation(self):
  def test_cuda_lu_lapack_getrf(self, dtype_name:str):
  if not config.enable_x64.value and dtype_name in ["f64", "c128"]:
  self.skipTest("Test disabled for x32 mode")
- if jaxlib_version < (0, 4, 32):
- self.skipTest("Not implemented in older versions of jaxlib")
  dtype = dict(f32=np.float32, f64=np.float64,
  c64=np.complex64, c128=np.complex128)[dtype_name]
  shape = (3, 4)
@@ -416,15 +404,12 @@ def test_cpu_qr_lapack_geqrf(self, dtype_name="f32"):
  data = self.load_testdata(cpu_qr_lapack_geqrf.data_2023_03_17[dtype_name])
  rtol = dict(f32=1e-3, f64=1e-5, c64=1e-3, c128=1e-5)[dtype_name]
  self.run_one_test(func, data, rtol=rtol)
- # TODO(b/344892332): Remove the check after the compatibility period.
- has_xla_ffi_support = jaxlib_version >= (0, 4, 32)
- if has_xla_ffi_support:
- with config.export_ignore_forward_compatibility(True):
- # FFI Kernel test
- data = self.load_testdata(
- cpu_qr_lapack_geqrf.data_2024_08_22[dtype_name]
- )
- self.run_one_test(func, data, rtol=rtol)
+ with config.export_ignore_forward_compatibility(True):
+ # FFI Kernel test
+ data = self.load_testdata(
+ cpu_qr_lapack_geqrf.data_2024_08_22[dtype_name]
+ )
+ self.run_one_test(func, data, rtol=rtol)
 
  @parameterized.named_parameters(
  dict(testcase_name=f"_dtype={dtype_name}_{batched}",
@@ -502,14 +487,11 @@ def test_cpu_lu_lapack_getrf(self, dtype_name:str):
  self.run_one_test(func, data, rtol=rtol, atol=atol,
  check_results=partial(self.check_lu_results, operand,
  dtype=dtype))
- # TODO(b/344892332): Remove the check after the compatibility period.
- has_xla_ffi_support = jaxlib_version >= (0, 4, 32)
- if has_xla_ffi_support:
- with config.export_ignore_forward_compatibility(True):
- # FFI Kernel test
- data = self.load_testdata(cpu_lu_lapack_getrf.data_2024_05_31[dtype_name])
- self.run_one_test(func, data, rtol=rtol, atol=atol,
- check_results=partial(self.check_lu_results, operand,
+ with config.export_ignore_forward_compatibility(True):
+ # FFI Kernel test
+ data = self.load_testdata(cpu_lu_lapack_getrf.data_2024_05_31[dtype_name])
+ self.run_one_test(func, data, rtol=rtol, atol=atol,
+ check_results=partial(self.check_lu_results, operand,
  dtype=dtype))
 
  def check_svd_results(self, input, res_run, res_exp,
@@ -629,16 +611,13 @@ def func(input):
  self.run_one_test(func, data, rtol=rtol, atol=atol,
  check_results=partial(self.check_svd_results,
  input))
- # TODO(b/344892332): Remove the check after the compatibility period.
- has_xla_ffi_support = jaxlib_version >= (0, 4, 32)
- if has_xla_ffi_support:
- with config.export_ignore_forward_compatibility(True):
- # FFI Kernel test
- data = self.load_testdata(
- cpu_svd_lapack_gesdd.data_2024_08_13[dtype_name]
- )
- self.run_one_test(func, data, rtol=rtol, atol=atol,
- check_results=partial(self.check_svd_results, input))
+ with config.export_ignore_forward_compatibility(True):
+ # FFI Kernel test
+ data = self.load_testdata(
+ cpu_svd_lapack_gesdd.data_2024_08_13[dtype_name]
+ )
+ self.run_one_test(func, data, rtol=rtol, atol=atol,
+ check_results=partial(self.check_svd_results, input))
 
  @jtu.parameterized_filterable(
  kwargs=[

tests/layout_test.py

@@ -15,7 +15,6 @@
 import contextlib
 import math
 from functools import partial
-import unittest
 from absl.testing import absltest
 import numpy as np
 
@@ -511,8 +510,6 @@ def g(x):
  'Layout passed to jit does not match the layout on the respective arg'):
  g(arr)
 
- @unittest.skipIf(xla_extension_version < 282,
- "Requires xla_extension_version >= 282")
  def test_in_layouts_jit_jnp_input(self):
  major_last_layout = DLL(major_to_minor=(1, 0))
  sharding = jax.sharding.SingleDeviceSharding(jax.devices()[0])

tests/linalg_test.py

@@ -16,7 +16,6 @@
 
 from functools import partial
 import itertools
-import unittest
 
 import numpy as np
 import scipy
@@ -2194,9 +2193,6 @@ def testHilbert(self, n):
  symmetrize_output=[True, False],
  )
  @jtu.skip_on_devices("tpu")
- @unittest.skipIf(
- jax._src.lib.version < (0, 4, 32), "requires jaxlib >= 0.4.32"
- )
  def testSymmetricProduct(self, shape, dtype, symmetrize_output):
  rng = jtu.rand_default(self.rng())
  batch_size = 10

tests/pjit_test.py

@@ -56,7 +56,6 @@
 from jax._src.lib.mlir import dialects
 from jax._src import xla_bridge
 from jax._src.lib import xla_client as xc
-from jax._src.lib import xla_extension_version
 from jax._src.lib import xla_extension
 from jax._src.util import curry, unzip2
 
@@ -4433,8 +4432,6 @@ def f(x):
  "Compiled object called with input sharding.*does not match"):
  compiled(cpu_arr)
 
- @unittest.skipIf(xla_extension_version < 281,
- 'Requires xla_extension_version >= 281')
  def test_different_devices_wsc_abstract_mesh_cache_hit(self):
  if jax.device_count() < 4:
  self.skipTest('Requires >=4 devices')
@@ -4463,8 +4460,6 @@ def f(x):
  self.assertEqual(lowering_count[0], 1)
  self.assertEqual(compilation_count[0], 2) # 2 misses since devices differ.
 
- @unittest.skipIf(xla_extension_version < 281,
- 'Requires xla_extension_version >= 281')
  def test_wsc_abstract_mesh(self):
  mesh = jtu.create_mesh((2, 2), ('x', 'y'))
  np_inp = np.arange(16).reshape(8, 2)
@@ -4484,8 +4479,6 @@ def f(x):
  self.assertArraysEqual(out_eager, np_inp * 2)
  self.assertEqual(out_eager.sharding, NamedSharding(mesh, P('x')))
 
- @unittest.skipIf(xla_extension_version < 281,
- 'Requires xla_extension_version >= 281')
  def test_wsc_sds_abstract_mesh(self):
  mesh = jtu.create_mesh((2,), 'x')
  s = NamedSharding(mesh, P())
@@ -4499,8 +4492,6 @@ def f(x):
  sds = jax.ShapeDtypeStruct((8, 2), np.float32, sharding=s)
  f.eval_shape(sds) # doesn't crash
 
- @unittest.skipIf(xla_extension_version < 281,
- 'Requires xla_extension_version >= 281')
  def test_wsc_vmap_abstract_mesh(self):
  mesh = jtu.create_mesh((2, 2), ('x', 'y'))
  s = NamedSharding(mesh, P('x', 'y'))
@@ -4517,8 +4508,6 @@ def f(x):
  out2 = jax.jit(jax.vmap(f, spmd_axis_name='y'))(arr)
  self.assertEqual(out2.sharding, NamedSharding(mesh, P('y', 'x')))
 
- @unittest.skipIf(xla_extension_version < 281,
- 'Requires xla_extension_version >= 281')
  def test_wsc_abstract_mesh_errors(self):
  mesh = jtu.create_mesh((2,), ('x',))
  np_inp = np.arange(8)

tests/shape_poly_test.py

@@ -2843,11 +2843,6 @@ def test_vmap_error(self):
  ((2, 3, 8, 4), "b1, b2, ..."),
  ((2, 3, 4, 5), "b1, b2, m, n"),
  ]
- # TODO(danfm): Remove once jaxlib v0.4.32 is the minimum version.
- # jaxlib versions before 0.4.32 require a static shape for the non-batch
- # dimensions because these are used for computing the "permuation_size"
- # which is passed to lu_pivots_to_permutation.
- if jaxlib_version >= (0, 4, 32) or not poly.endswith("m, n")
  ],
  [
  # The random primitive tests, with threefry (both partitionable and