Repository: google/jax
Pull request #22635 (2024-07-24 17:31:26)

Remove dead code now that xmap is deleted
Remove dead code now that xmap is deleted



Patch:
diff --git a/jax/_src/interpreters/pxla.py b/jax/_src/interpreters/pxla.py
index b05ae2d0a..886b7eec7 100644
--- a/jax/_src/interpreters/pxla.py
+++ b/jax/_src/interpreters/pxla.py
@@ -31,7 +31,6 @@ import warnings
 import numpy as np
 
 import jax
-from jax.errors import JAXTypeError
 
 from jax._src import api_util
 from jax._src import compiler
@@ -1505,8 +1504,6 @@ def _pmap_lowering(ctx, *in_nodes, axis_name,
 mlir.register_lowering(xla_pmap_p, _pmap_lowering)
 
 
-# ------------------- xmap -------------------
-
 def tile_aval_nd(axis_sizes, in_axes: ArrayMapping, aval):
   assert isinstance(aval, ShapedArray)
   shape = list(aval.shape)
@@ -1537,34 +1534,6 @@ def mesh_global_to_local(mesh, axes: ArrayMapping, aval):
                         tile_aval_nd(mesh.shape, axes, aval))
 
 
-class SPMDBatchTrace(batching.BatchTrace):
-  def get_axis_primitive_batcher(self, primitive, frame):
-    if primitive in spmd_primitive_batchers:
-      return partial(spmd_primitive_batchers[primitive],
-          frame.size, frame.name, frame.main_trace.trace_type)
-    return super().get_axis_primitive_batcher(primitive, frame)
-
-
-spmd_primitive_batchers: dict[core.Primitive, Callable] = {}
-
-
-def vtile_by_mesh(fun: lu.WrappedFun,
-                  mesh: Mesh,
-                  in_axes: Sequence[ArrayMapping],
-                  out_axes: Sequence[ArrayMapping]):
-  # We vectorize in reversed order, because vmap is often biased towards
-  # moving the batch axis to the front, and this way of stacking transforms
-  # will order the batch axes according to the mesh axis order.
-  # Not strictly necessary, but seems nicer than reversing it?
-  for name, size in reversed(mesh.shape.items()):
-    fun = batching.vtile(fun,
-                         tuple(a.get(name, None) for a in in_axes),
-                         tuple(a.get(name, None) for a in out_axes),
-                         tile_size=size,
-                         axis_name=name,
-                         main_type=SPMDBatchTrace)
-  return fun
-
 full_to_shard_p = core.Primitive('full_to_shard')
 
 @full_to_shard_p.def_abstract_eval
@@ -1638,30 +1607,6 @@ def _shard_to_full_lowering(ctx: mlir.LoweringRuleContext, x, *, axes: ArrayMapp
   return (mlir.wrap_with_shard_to_full_op(ctx, sx, aval_out, sharding_proto,
                                           unspecified_dims),)
 
-@lu.transformation
-def vtile_manual(manual_axes: frozenset[sharding_impls.MeshAxisName],
-                 mesh: Mesh,
-                 in_axes: Sequence[ArrayMapping],
-                 out_axes: Sequence[ArrayMapping],
-                 *args):
-  tiled_args = [full_to_shard_p.bind(arg, axes=axes, mesh=mesh, manual_axes=manual_axes)
-                for arg, axes in zip(args, in_axes)]
-  tiled_outs = yield tiled_args, {}
-  outs = [shard_to_full_p.bind(out, axes=axes, mesh=mesh, manual_axes=manual_axes)
-          for out, axes in zip(tiled_outs, out_axes)]
-  yield outs
-
-
-@dataclasses.dataclass(frozen=True)
-class TileVectorize:
-  pass
-
-@dataclasses.dataclass(frozen=True)
-class TileManual:
-  manual_axes: frozenset[sharding_impls.MeshAxisName]
-
-TilingMethod = Union[TileVectorize, TileManual]
-
 
 def check_if_any_auto(
     shardings: Iterable[(JSharding | AUTO | UnspecifiedValue)]) -> bool:
@@ -2333,171 +2278,6 @@ def _to_logical_sharding(
     raise TypeError(aval)
 
 
-@profiler.annotate_function
-def lower_mesh_computation(
-    fun_or_jaxpr: lu.WrappedFun | core.ClosedJaxpr,
-    api_name: str,
-    fun_name: str,
-    mesh: Mesh,
-    in_shardings: Sequence[sharding_impls.NamedSharding | AUTO],
-    out_shardings: Sequence[(sharding_impls.NamedSharding | AUTO |
-                                  UnspecifiedValue)],
-    donated_invars: Sequence[bool],
-    spmd_lowering: bool,
-    global_in_avals: Sequence[core.ShapedArray],
-    tiling_method: TilingMethod | None,
-    lowering_platforms: tuple[str, ...] | None,
-    lowering_parameters: mlir.LoweringParameters) -> MeshComputation:
-  assert not mesh.empty
-  backend = xb.get_device_backend(mesh.devices.flat[0])
-  platforms = lowering_platforms or (backend.platform,)
-  name_stack = source_info_util.new_name_stack(wrap_name(fun_name, api_name))
-
-  global_axis_sizes = mesh.shape
-
-  log_priority = logging.WARNING if config.log_compiles.value else logging.DEBUG
-  if logger.isEnabledFor(log_priority):
-    logger.log(log_priority,
-               "Compiling %s for %s mesh with global shapes and types %s. "
-               "Argument mapping: %s.",
-               fun_name, tuple(global_axis_sizes.items()), global_in_avals,
-               in_shardings)
-
-  # 1. Trace to jaxpr and preprocess/verify it
-  if spmd_lowering:
-    manual_axes: frozenset[MeshAxisName] = frozenset()
-    # TODO: Consider handling xmap's 'vectorize' in here. We can vmap once instead of vtile twice!
-    if tiling_method is not None:
-      if isinstance(tiling_method, TileVectorize):
-        tiling_transform = vtile_by_mesh
-      elif isinstance(tiling_method, TileManual):
-        tiling_transform = lambda f, *args: vtile_manual(f, tiling_method.manual_axes, *args)
-        manual_axes = tiling_method.manual_axes
-      else:
-        raise NotImplementedError(f"Unrecognized tiling method: {tiling_method}")
-      assert not callable(out_shardings)
-      assert isinstance(fun_or_jaxpr, lu.WrappedFun)
-      # This is the xmap path where there is no `AUTO` or `UNSPECIFIED`, which
-      # is why `.spec` can be accessed.
-      fun_or_jaxpr = tiling_transform(
-          fun_or_jaxpr, mesh, [get_array_mapping(i.spec) for i in in_shardings],  # type: ignore
-          [get_array_mapping(o.spec) for o in out_shardings])  # type: ignore
-    in_jaxpr_avals = global_in_avals
-  else:
-    assert isinstance(tiling_method, TileVectorize)
-    # In non-spmd lowering path, there is no `AUTO` or `UNSPECIFIED`, which is
-    # why `.spec` can be accessed.
-    in_tiled_avals = [tile_aval_nd(global_axis_sizes, get_array_mapping(i.spec), aval)  # type: ignore
-                      for aval, i in safe_zip(global_in_avals, in_shardings)]
-    in_jaxpr_avals = in_tiled_avals
-
-  with core.extend_axis_env_nd(mesh.shape.items()):
-    if isinstance(fun_or_jaxpr, lu.WrappedFun):
-      with dispatch.log_elapsed_time(
-          "Finished tracing + transforming {fun_name} in {elapsed_time} sec",
-          fun_name=str(name_stack), event=dispatch.JAXPR_TRACE_EVENT):
-        jaxpr, out_jaxpr_avals, consts = pe.trace_to_jaxpr_final(
-            fun_or_jaxpr, in_jaxpr_avals)
-    else:
-      assert isinstance(fun_or_jaxpr, core.ClosedJaxpr)
-      jaxpr = fun_or_jaxpr.jaxpr
-      out_jaxpr_avals = fun_or_jaxpr.out_avals
-      consts = fun_or_jaxpr.consts
-
-  assert len(out_shardings) == len(out_jaxpr_avals)
-  if spmd_lowering:
-    global_out_avals = out_jaxpr_avals
-  else:
-    # In non-spmd lowering path, there is no `AUTO` or `UNSPECIFIED`, which is
-    # why `.spec` can be accessed.
-    global_out_avals = [untile_aval_nd(global_axis_sizes, get_array_mapping(o.spec), aval)  # type: ignore
-                        for aval, o in safe_zip(out_jaxpr_avals, out_shardings)]
-
-  _sanitize_mesh_jaxpr(jaxpr)
-  jaxpr = dispatch.apply_outfeed_rewriter(jaxpr)
-
-  # 2. Build up the HLO
-  tuple_args = dispatch.should_tuple_args(len(in_jaxpr_avals), backend.platform)
-
-  in_partitions: list[JSharding | None] | None
-  out_partitions: list[JSharding | None] | None
-  axis_ctx: mlir.AxisContext
-  if spmd_lowering:
-    in_partitions = map(_to_logical_sharding, global_in_avals, in_shardings)
-    out_partitions = map(_to_logical_sharding, global_out_avals, out_shardings)
-    replicated_args = [False] * len(in_jaxpr_avals)
-    axis_ctx = sharding_impls.SPMDAxisContext(mesh, manual_axes)
-    num_replicas = 1
-    num_partitions = mesh.devices.size
-  else:
-    replicated_args = [not get_array_mapping(i.spec) for i in in_shardings]  # type: ignore
-    in_partitions = None
-    out_partitions = None
-    axis_env = sharding_impls.AxisEnv(
-        nreps=mesh.size,
-        names=tuple(global_axis_sizes.keys()),
-        sizes=tuple(global_axis_sizes.values()))
-    axis_ctx = sharding_impls.ReplicaAxisContext(axis_env)
-    num_replicas = mesh.devices.size
-    num_partitions = 1
-  jaxpr = core.remove_named_axis_effects(jaxpr, mesh.axis_names)
-  closed_jaxpr = core.ClosedJaxpr(jaxpr, consts)
-  module_name = f"{api_name}_{fun_name}"
-  with core.extend_axis_env_nd(mesh.shape.items()):
-    if any(effects.ordered_effects.contains(eff) for eff
-           in closed_jaxpr.effects):
-      raise ValueError("Ordered effects not supported in mesh computations.")
-    unordered_effects = list(effects.ordered_effects.filter_not_in(
-      closed_jaxpr.effects))
-    ordered_effects = list(effects.ordered_effects.filter_in(
-      closed_jaxpr.effects))
-    with dispatch.log_elapsed_time(
-        "Finished jaxpr to MLIR module conversion {fun_name} in {elapsed_time} sec",
-        fun_name=str(name_stack), event=dispatch.JAXPR_TO_MLIR_MODULE_EVENT):
-      lowering_result = mlir.lower_jaxpr_to_module(
-          module_name,
-          closed_jaxpr,
-          ordered_effects=ordered_effects,
-          backend_or_name=backend,
-          platforms=platforms,
-          axis_context=axis_ctx,
-          name_stack=name_stack,
-          donated_args=donated_invars,
-          replicated_args=replicated_args,
-          arg_shardings=in_partitions,
-          result_shardings=out_partitions,
-          arg_names=jaxpr.debug_info and jaxpr.debug_info.arg_names,
-          result_names=jaxpr.debug_info and jaxpr.debug_info.result_paths,
-          num_replicas=num_replicas,
-          num_partitions=num_partitions,
-          lowering_parameters=lowering_parameters)
-
-  return MeshComputation(
-      str(name_stack),
-      lowering_result.module,
-      donated_invars,
-      platforms,
-      global_in_avals=global_in_avals,
-      global_out_avals=global_out_avals,
-      in_shardings=in_shardings,
-      out_shardings=out_shardings,
-      spmd_lowering=spmd_lowering,
-      tuple_args=tuple_args,
-      auto_spmd_lowering=False,
-      unordered_effects=unordered_effects,
-      ordered_effects=ordered_effects,
-      host_callbacks=lowering_result.host_callbacks,
-      keepalive=lowering_result.keepalive,
-      kept_var_idx=set(range(len(global_in_avals))),
-      backend=backend,
-      device_assignment=_create_da_object(tuple(mesh.devices.flat)),
-      committed=True,
-      in_layouts=(None,) * len(global_in_avals),
-      out_layouts=(None,) * len(global_out_avals),
-      shape_poly_state=lowering_result.shape_poly_state,
-      all_args_info=None,
-      pgle_profiler=None)
-
 class MeshComputation(stages.XlaLowering):
   _hlo: ir.Module
   _executable: MeshExecutable | None
@@ -3355,64 +3135,6 @@ def get_array_mapping(pspec: PartitionSpec) -> ArrayMappingOrAutoOrUnspecified:
   return _get_array_mapping(parsed_pspec)
 
 
-_forbidden_primitives = {
-  'xla_pmap': 'pmap',
-}
-def _sanitize_mesh_jaxpr(jaxpr):
-  if isinstance(jaxpr, core.ClosedJaxpr):
-    jaxpr = jaxpr.jaxpr
-  for eqn in jaxpr.eqns:
-    if eqn.primitive.name in _forbidden_primitives:
-      raise RuntimeError(f"Nesting {_forbidden_primitives[eqn.primitive.name]} "
-                         f"inside xmaps not supported!")
-    core.traverse_jaxpr_params(_sanitize_mesh_jaxpr, eqn.params)
-
-
-custom_resource_typing_rules: dict[core.Primitive, Callable] = {}
-
-def resource_typecheck(jaxpr, resource_env, axis_resources, what_jaxpr_thunk):
-  if isinstance(jaxpr, core.ClosedJaxpr):
-    jaxpr = jaxpr.jaxpr
-  def _check_aval(aval, what_thunk):
-    if not hasattr(aval, 'named_shape'):
-      return
-    resource_to_axis = {}
-    for axis in aval.named_shape:
-      if axis_resources:
-        for resource in axis_resources[axis]:
-          if resource in resource_to_axis:
-            other_axis = resource_to_axis[resource]
-            axis, other_axis = sorted([str(axis), str(other_axis)])
-            raise JAXTypeError(
-                f"Axes `{axis}` and `{other_axis}` are both mapped to the "
-                f"resource `{resource}`, but they coincide in the named_shape "
-                f"of {what_thunk()}")
-          resource_to_axis[resource] = axis
-
-  what_thunk = lambda: (f"an input to {what_jaxpr_thunk()}")
-  for v in jaxpr.constvars:
-    _check_aval(v.aval, what_thunk)
-  for v in jaxpr.invars:
-    _check_aval(v.aval, what_thunk)
-  what_thunk = lambda: (f"a value returned from a primitive {eqn.primitive} created "
-                        f"at {source_info_util.summarize(eqn.source_info)}")
-  rec_what_jaxpr_thunk = lambda: (f"a primitive {eqn.primitive} created at"
-                                  f"{source_info_util.summarize(eqn.source_info)}")
-  for eqn in jaxpr.eqns:
-    typing_rule = custom_resource_typing_rules.get(eqn.primitive, None)
-    if typing_rule:
-      typing_rule([v.aval for v in eqn.invars], eqn.params, eqn.source_info,
-                  resource_env, axis_resources)
-    else:
-      core.traverse_jaxpr_params(partial(resource_typecheck,
-                                         resource_env=resource_env,
-                                         axis_resources=axis_resources,
-                                         what_jaxpr_thunk=rec_what_jaxpr_thunk),
-                                 eqn.params)
-    for v in eqn.outvars:
-      _check_aval(v.aval, what_thunk)
-
-
 @contextmanager
 def maybe_extend_axis_env(*args, **kwargs):
   with core.extend_axis_env(*args, **kwargs):
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index e2fc7c8d3..a2b5993ea 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -52,7 +52,6 @@ from jax._src.api_util import (
     donation_vector, shaped_abstractify, check_callable, resolve_argnums,
     argnames_partial_except, debug_info, result_paths, jaxpr_debug_info,
     hoist_obj_attrs)
-from jax._src.errors import JAXTypeError
 from jax._src.interpreters import partial_eval as pe
 from jax._src.partition_spec import PartitionSpec
 from jax._src.interpreters import xla
@@ -69,7 +68,7 @@ from jax._src import sharding
 from jax._src.sharding_impls import (
     NamedSharding, GSPMDSharding,
     SingleDeviceSharding, PmapSharding, AUTO, UNSPECIFIED, UnspecifiedValue,
-    ParsedPartitionSpec, SpecSync, get_single_pspec, is_auto, is_unspecified,
+    ParsedPartitionSpec, get_single_pspec, is_unspecified,
     is_unspecified_or_auto, prepare_axis_resources, parse_flatten_op_sharding)
 from jax._src.layout import Layout, DeviceLocalLayout, AutoLayout
 from jax._src.state import discharge as state_discharge, RefEffect, AbstractRef
@@ -1757,9 +1756,6 @@ def _pjit_lower_cached(
     lowering_platforms: tuple[str, ...] | None,
     lowering_parameters: mlir.LoweringParameters,
     pgle_profiler: profiler.PGLEProfiler | None):
-  if resource_env is not None:
-    pxla.resource_typecheck(jaxpr, resource_env, {}, lambda: "pjit")
-
   if resource_env is not None:
     mesh = resource_env.physical_mesh
     api_name = 'pjit'
@@ -1768,25 +1764,15 @@ def _pjit_lower_cached(
     mesh = None
     api_name = 'jit'
 
-  # For `pjit(xmap)` cases, it needs to take the `lower_mesh_computation` path
-  # because `xmap` only supports SPMDAxisContext right now.
-  if dispatch.jaxpr_has_primitive(jaxpr.jaxpr, 'xmap'):
-    return pxla.lower_mesh_computation(
-      jaxpr, api_name, name, mesh,
-      in_shardings, out_shardings, donated_invars,
-      True, jaxpr.in_avals, tiling_method=None,
+  return pxla.lower_sharding_computation(
+      jaxpr, api_name, name, in_shardings, out_shardings,
+      in_layouts, out_layouts, tuple(donated_invars),
+      keep_unused=keep_unused,
+      devices_from_context=(
+          None if mesh is None or mesh.empty else list(mesh.devices.flat)),
       lowering_platforms=lowering_platforms,
-      lowering_parameters=lowering_parameters)
-  else:
-    return pxla.lower_sharding_computation(
-        jaxpr, api_name, name, in_shardings, out_shardings,
-        in_layouts, out_layouts, tuple(donated_invars),
-        keep_unused=keep_unused,
-        devices_from_context=(
-            None if mesh is None or mesh.empty else list(mesh.devices.flat)),
-        lowering_platforms=lowering_platforms,
-        lowering_parameters=lowering_parameters,
-        pgle_profiler=pgle_profiler)
+      lowering_parameters=lowering_parameters,
+      pgle_profiler=pgle_profiler)
 
 
 def pjit_staging_rule(trace, *args, **params):
@@ -2007,7 +1993,6 @@ def _pjit_batcher(insert_axis, spmd_axis_name,
 
 batching.spmd_axis_primitive_batchers[pjit_p] = partial(_pjit_batcher, False)
 batching.axis_primitive_batchers[pjit_p] = partial(_pjit_batcher, False, None)
-pxla.spmd_primitive_batchers[pjit_p] = partial(_pjit_batcher, True, None)
 
 def _pjit_batcher_for_sharding(
     s: sharding.Sharding | UnspecifiedValue,
@@ -2369,66 +2354,6 @@ def dce_jaxpr_pjit_rule(used_outputs: list[bool], eqn: core.JaxprEqn
 pe.dce_rules[pjit_p] = dce_jaxpr_pjit_rule
 
 
-def _check_resources_against_named_axes(what, aval, pos_axis_resources, named_axis_resources):
-  pjit_resources = set(
-      it.chain.from_iterable([d for d in pos_axis_resources if d is not None]))
-  aval_resources = set(it.chain.from_iterable(
-    named_axis_resources[a] for a in aval.named_shape))
-  overlap = pjit_resources & aval_resources
-  if overlap:
-    raise JAXTypeError(
-        f"{what} has an axis resources specification of "
-        f"{pos_axis_resources.unsynced_user_spec(SpecSync.DIM_PERMUTE)} "
-        f"that uses one or more mesh axes already used by xmap to partition "
-        f"a named axis appearing in its named_shape (both use mesh axes "
-        f"{mesh_lib.show_axes(overlap)})")
-
-def _resource_typing_pjit(avals, params, source_info, resource_env, named_axis_resources):
-  jaxpr = params["jaxpr"]
-  what = "pjit input"
-  if (resource_env is not None and params['resource_env'] is not None and
-      resource_env.physical_mesh != params['resource_env'].physical_mesh):
-      raise RuntimeError("Changing the physical mesh is not allowed inside pjit.")
-
-  for aval, s in zip(jaxpr.in_avals, params['in_shardings']):
-    if is_unspecified(s) or is_auto(s):
-      continue
-    elif hasattr(s, '_parsed_pspec'):
-      parsed_pspec = s._parsed_pspec
-    else:
-      if resource_env is not None and not resource_env.physical_mesh.empty:
-        parsed_pspec = parse_flatten_op_sharding(
-            s._to_xla_hlo_sharding(aval.ndim), resource_env.physical_mesh)[0]
-      else:
-        parsed_pspec = None
-    if parsed_pspec is not None:
-      _check_resources_against_named_axes(what, aval, parsed_pspec,
-                                          named_axis_resources)
-
-  pxla.resource_typecheck(
-      jaxpr.jaxpr, resource_env, named_axis_resources,
-      lambda: (f"a pjit'ed function {params['name']} "
-               f"(pjit called at {source_info_util.summarize(source_info)})"))
-
-  what = "pjit output"
-  for aval, s in zip(jaxpr.out_avals, params['out_shardings']):
-    if is_unspecified(s) or is_auto(s):
-      continue
-    elif hasattr(s, '_parsed_pspec'):
-      parsed_pspec = s._parsed_pspec
-    else:
-      if resource_env is not None and not resource_env.physical_mesh.empty:
-        parsed_pspec = parse_flatten_op_sharding(
-            s._to_xla_hlo_sharding(aval.ndim), resource_env.physical_mesh)[0]
-      else:
-        parsed_pspec = None
-    if parsed_pspec is not None:
-      _check_resources_against_named_axes(what, aval, parsed_pspec,
-                                          named_axis_resources)
-
-pxla.custom_resource_typing_rules[pjit_p] = _resource_typing_pjit
-
-
 def _pjit_pp_rule(eqn, context, settings):
   params = dict(eqn.params)
   del params['inline']
@@ -2637,27 +2562,6 @@ batching.spmd_axis_primitive_batchers[sharding_constraint_p] = partial(
     _sharding_constraint_batcher, False)
 batching.axis_primitive_batchers[sharding_constraint_p] = partial(
     _sharding_constraint_batcher, False, None)
-pxla.spmd_primitive_batchers[sharding_constraint_p] = partial(
-    _sharding_constraint_batcher, True, None)
-
-
-def _resource_typing_sharding_constraint(avals, params, source_info,
-                                         resource_env, named_axis_resources):
-  aval, = avals
-  parsed_pspec = None
-  if isinstance(params['sharding'], NamedSharding):
-    parsed_pspec = params['sharding']._parsed_pspec
-  else:
-    if not resource_env.physical_mesh.empty:
-      parsed_pspec = parse_flatten_op_sharding(
-          params['sharding']._to_xla_hlo_sharding(aval.ndim),
-          resource_env.physical_mesh)[0]
-  if parsed_pspec is not None:
-    _check_resources_against_named_axes(
-      "with_sharding_constraint input", aval, parsed_pspec, named_axis_resources)
-
-pxla.custom_resource_typing_rules[sharding_constraint_p] = \
-    _resource_typing_sharding_constraint
 
 # -------------------- helpers --------------------
 


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 12201a3b0..9ed2405bc 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4499,8 +4499,8 @@ class PJitErrorTest(jtu.JaxTestCase):
         return h(x)
     xshape = (2, 5, 6)
     x = jnp.arange(math.prod(xshape)).reshape(xshape)
-    with self.assertRaisesRegex(RuntimeError,
-                                "Changing the physical mesh is not allowed.*"):
+    with self.assertRaisesRegex(
+        ValueError, "Received incompatible devices for pjitted computation.*"):
       f(x)
 
   @parameterized.named_parameters(