Remove code to support jaxlib < 0.5.3.

The new xla_extension_version is 320.

PiperOrigin-RevId: 738522486

Author: hawkinsp

jax/_src/export/_export.py

@@ -43,7 +43,7 @@
 from jax._src.interpreters import mlir
 from jax._src.interpreters import pxla
 from jax._src.lib import xla_client
-from jax._src.lib import xla_extension, xla_extension_version
+from jax._src.lib import xla_extension
 from jax._src.lib.mlir import ir, passmanager
 from jax._src.lib.mlir.dialects import hlo
 from jax._src.lib.mlir.dialects import func as func_dialect
@@ -674,10 +674,8 @@ def _export_lowered(
   # Shardy was used during lowering if we can find the Shardy mesh in the
   # module. Note that the mesh should have been lifted by the
   # `sdy-lift-inlined-meshes` pass in mlir.py.
-  shardy_enabled = False
-  if xla_extension_version >= 319:
-    shardy_enabled = xla_extension.sdy.lowered_with_shardy(
-        mlir.module_to_bytecode(mlir_module))
+  shardy_enabled = xla_extension.sdy.lowered_with_shardy(
+      mlir.module_to_bytecode(mlir_module))
 
   mlir_module_serialized = _module_to_bytecode(mlir_module, shardy_enabled)
 
@@ -784,7 +782,7 @@ def _get_exported_vjp(exp_primal: Exported) -> Exported:
       _get_vjp=_get_exported_vjp)
 
 def _module_to_bytecode(module: ir.Module, shardy_enabled: bool) -> bytes:
-  if xla_extension_version >= 319 and shardy_enabled:
+  if shardy_enabled:
     mlir_str = xla_extension.sdy.sdy_round_trip_export_pipeline(
         mlir.module_to_bytecode(module))
   else:
@@ -1423,10 +1421,8 @@ def _call_exported_lowering(ctx: mlir.LoweringRuleContext, *args,
     ctx.module_context.shape_poly_state.uses_dim_vars = True
   submodule = ir.Module.parse(exported.mlir_module())
 
-  shardy_enabled = False
-  if xla_extension_version >= 319:
-    shardy_enabled = xla_extension.sdy.lowered_with_shardy(
-        mlir.module_to_bytecode(submodule))
+  shardy_enabled = xla_extension.sdy.lowered_with_shardy(
+      mlir.module_to_bytecode(submodule))
   if shardy_enabled:
     submodule = ir.Module.parse(xla_extension.sdy.sdy_round_trip_import_shardings(
         mlir.module_to_bytecode(submodule)))

jax/_src/interpreters/mlir.py

@@ -55,7 +55,7 @@
                                      SdyArraySharding, SdyArrayShardingList)
 from jax._src.util import foreach
 from jax._src.lib import xla_client as xc
-from jax._src.lib import xla_extension, xla_extension_version
+from jax._src.lib import xla_extension
 from jax._src.lib.mlir import dialects, ir, passmanager
 from jax._src.lib.mlir.dialects import func as func_dialect, hlo
 from jax._src.lib.mlir import register_jax_dialects
@@ -3031,11 +3031,8 @@ def refine_polymorphic_shapes(module: ir.Module) -> ir.Module:
             mlir_module=module_to_bytecode(module),
             enable_shape_assertions=True,
             validate_static_shapes=True)
-    if xla_extension_version >= 319:
-      refined_module_str = refine_polymorphic_shapes(
-          enable_shardy=config.use_shardy_partitioner.value)
-    else:
-      refined_module_str = refine_polymorphic_shapes()
+    refined_module_str = refine_polymorphic_shapes(
+        enable_shardy=config.use_shardy_partitioner.value)
   except Exception as e:
     raise ValueError(
         "Error refining shapes. " +

jax/_src/lax/lax.py

@@ -66,7 +66,6 @@
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import chlo
 from jax._src.lib.mlir.dialects import hlo
-from jax._src.lib import xla_extension_version
 from jax._src.sharding_impls import (PmapSharding, NamedSharding,
                                      PartitionSpec as P, canonicalize_sharding)
 from jax._src.typing import Array, ArrayLike, DimSize, DuckTypedArray, DTypeLike, Shape
@@ -2267,11 +2266,6 @@ def _convert_to_hlo_attr(self, lhs_dtype: DTypeLike,
       case DotAlgorithmPreset.BF16_BF16_F32_X6:
         return hlo.DotAlgorithm.get(bf16, bf16, f32, 1, 1, 6, False)
       case DotAlgorithmPreset.BF16_BF16_F32_X9:
-        if xla_extension_version < 320:
-          raise ValueError(
-              "The dot algorithm BF16_BF16_F32_X9 requires XLA extension "
-              "version >= 320."
-          )
         return hlo.DotAlgorithm.get(bf16, bf16, f32, 1, 1, 9, False)
       case DotAlgorithmPreset.TF32_TF32_F32:
         return hlo.DotAlgorithm.get(tf32, tf32, f32, 1, 1, 1, False)

jax/_src/sharding_impls.py

@@ -33,7 +33,6 @@
 from jax._src import xla_bridge as xb
 from jax._src import mesh_utils
 from jax._src.lib import xla_client as xc
-from jax._src.lib import xla_extension_version
 from jax._src.lib.mlir.dialects import sdy
 from jax._src.named_sharding import (  # noqa: F401
     SdyArraySharding, SdyDimSharding, UnspecifiedValue, AUTO,
@@ -881,8 +880,7 @@ def parse_flatten_op_sharding(
     return out
   elif hlo_sharding.is_replicated():
     return [PartitionSpec()]
-  elif (xla_extension_version >= 319 and hlo_sharding.is_maximal()
-        and mesh.size == 1):
+  elif hlo_sharding.is_maximal() and mesh.size == 1:
     return [PartitionSpec()]
   elif hlo_sharding.is_tiled():
     mesh_shape = mesh.shape

jax/_src/util.py

@@ -108,11 +108,7 @@ def foreach(f, *args):
     return None
 
 else:
-  # TODO(phawkins): remove after jaxlib 0.5.2 is the minimum.
-  if hasattr(jaxlib_utils, 'foreach'):
-    foreach = jaxlib_utils.foreach
-  else:
-    foreach = safe_map
+  foreach = jaxlib_utils.foreach
 
 
 def unzip2(xys: Iterable[tuple[T1, T2]]
@@ -244,61 +240,8 @@ def curry(f):
   """
   return wraps(f)(partial(partial, f))
 
-# TODO(phawkins): make this unconditional after jaxlib 0.5.3 is the minimum.
 toposort: Callable[[Iterable[Any]], list[Any]]
-if hasattr(jaxlib_utils, "topological_sort"):
-  toposort = partial(jaxlib_utils.topological_sort, "parents")
-else:
-
-  def toposort(end_nodes):
-    if not end_nodes:
-      return []
-    end_nodes = _remove_duplicates(end_nodes)
-
-    child_counts = {}
-    stack = list(end_nodes)
-    while stack:
-      node = stack.pop()
-      if id(node) in child_counts:
-        child_counts[id(node)] += 1
-      else:
-        child_counts[id(node)] = 1
-        stack.extend(node.parents)
-    for node in end_nodes:
-      child_counts[id(node)] -= 1
-
-    sorted_nodes = []
-    childless_nodes = [
-        node for node in end_nodes if child_counts[id(node)] == 0
-    ]
-    assert childless_nodes
-    while childless_nodes:
-      node = childless_nodes.pop()
-      sorted_nodes.append(node)
-      for parent in node.parents:
-        if child_counts[id(parent)] == 1:
-          childless_nodes.append(parent)
-        else:
-          child_counts[id(parent)] -= 1
-    sorted_nodes = sorted_nodes[::-1]
-
-    check_toposort(sorted_nodes)
-    return sorted_nodes
-
-  def check_toposort(nodes):
-    visited = set()
-    for node in nodes:
-      assert all(id(parent) in visited for parent in node.parents)
-      visited.add(id(node))
-
-  def _remove_duplicates(node_list):
-    seen = set()
-    out = []
-    for n in node_list:
-      if id(n) not in seen:
-        seen.add(id(n))
-        out.append(n)
-    return out
+toposort = partial(jaxlib_utils.topological_sort, "parents")
 
 
 def split_merge(predicate, xs):
@@ -320,7 +263,6 @@ def merge(new_lhs, new_rhs):
 
   return lhs, rhs, merge
 
-
 def _ignore(): return None
 
 

jax/experimental/mosaic/gpu/dialect_lowering.py

@@ -849,13 +849,9 @@ def _mgpu_wait_op_lowering_rule(
   return []
 
 
-# TODO(bchetioui): remove this once jaxlib minimum version >= 0.5.2.
-SliceSMEMOp = getattr(mgpu, "SliceSMEMOp", None)
-
-
-@_register_lowering(SliceSMEMOp)
+@_register_lowering(mgpu.SliceSMEMOp)
 def _mgpu_slice_smem_op_lowering_rule(
-    ctx: LoweringContext, op: SliceSMEMOp
+    ctx: LoweringContext, op: mgpu.SliceSMEMOp
 ) -> Sequence[ir.Value]:
   del ctx
   return [_slice_smem(op.result.type, op.offset)]

jax/experimental/mosaic/gpu/transform_inference.py

@@ -172,11 +172,9 @@ def _infer_vector_load_store_transforms(
 
   return None
 
-# TODO(bchetioui): remove this once jaxlib minimum version >= 0.5.2.
-SliceSMEMOp = getattr(mgpu, "SliceSMEMOp", None)
 
-@partial(_add_transform_inference_rule, SliceSMEMOp)
-def _infer_slice_smem_transforms(op: SliceSMEMOp) -> OptionalTransforms:
+@partial(_add_transform_inference_rule, mgpu.SliceSMEMOp)
+def _infer_slice_smem_transforms(op: mgpu.SliceSMEMOp) -> OptionalTransforms:
   transforms = None
   uses = cast(ir.OpResult, op.result).uses
 

jax/experimental/sparse/linalg.py

@@ -29,7 +29,6 @@
 from jax._src import core
 from jax._src import ffi
 from jax._src.interpreters import ad
-from jax._src.lib import gpu_solver
 
 import numpy as np
 from scipy.sparse import csr_matrix, linalg
@@ -534,11 +533,6 @@ def _spsolve_abstract_eval(data, indices, indptr, b, *, tol, reorder):
 
 
 def _spsolve_gpu_lowering(ctx, data, indices, indptr, b, *, tol, reorder):
-  # TODO(danfm): remove after JAX 0.5.1 release.
-  if hasattr(gpu_solver, "cuda_csrlsvqr"):
-    data_aval, _, _, _, = ctx.avals_in
-    return gpu_solver.cuda_csrlsvqr(data_aval.dtype, data, indices,
-                                    indptr, b, tol, reorder)
   return ffi.ffi_lowering("cusolver_csrlsvqr_ffi")(
       ctx, data, indices, indptr, b, tol=np.float64(tol),
       reorder=np.int32(reorder))

tests/lax_test.py

@@ -49,7 +49,6 @@
 from jax._src.lax import lax as lax_internal
 from jax._src.util import NumpyComplexWarning, safe_zip
 from jax._src.tree_util import tree_map
-from jax._src.lib import xla_extension_version
 
 config.parse_flags_with_absl()
 
@@ -1128,11 +1127,6 @@ def testDotAlgorithm(self, algorithm, dtype):
         raise SkipTest(
             f"The dot algorithm '{algorithm}' is not supported on CPU.")
     if jtu.test_device_matches(["gpu"]):
-      if (algorithm == lax.DotAlgorithmPreset.BF16_BF16_F32_X9 and
-          xla_extension_version < 320):
-        raise SkipTest(
-              f"The dot algorithm ${algorithm} requires XLA extension version "
-              ">= 320.")
       # GPU algorithm support is a little spotty. It is checked in
       # xla/service/algorithm_util.cc and the logic is copied here.
       if algorithm in {

tests/linalg_test.py

@@ -867,9 +867,6 @@ def testSVD(self, b, m, n, dtype, full_matrices, compute_uv, hermitian, algorith
         self.skipTest("Hermitian SVD doesn't support the algorithm parameter.")
       if not jtu.test_device_matches(["cpu", "gpu"]):
         self.skipTest("SVD algorithm selection only supported on CPU and GPU.")
-      # TODO(danfm): Remove this check after 0.5.2 is released.
-      if jtu.test_device_matches(["cpu"]) and jtu.jaxlib_version() <= (0, 5, 1):
-        self.skipTest("SVD algorithm selection on CPU requires a newer jaxlib version.")
       if jtu.test_device_matches(["cpu"]) and algorithm == lax.linalg.SvdAlgorithm.JACOBI:
         self.skipTest("Jacobi SVD not supported on GPU.")