[sharding_in_types] Make argmax and argmin work with sharding_in_types. This also requires adding reduce_p sharding rule

PiperOrigin-RevId: 699244204

Author: yashk2810

jax/_src/lax/lax.py

@@ -5007,6 +5007,11 @@ def _reduce_shape_rule(*avals, computation, jaxpr, dimensions):
     raise ValueError(f'reduce found non-scalar initial value: {init_val_shapes}')
   return [tuple(np.delete(op.shape, dimensions)) for op in operand_avals]
 
+def _reduce_sharding_rule(*avals, computation, jaxpr, dimensions):
+  operand_avals, _ = split_list(avals, [len(avals) // 2])
+  return [op.sharding.with_spec(tuple_delete(op.sharding.spec, dimensions))
+          for op in operand_avals]
+
 def _reduce_dtype_rule(*avals, computation, jaxpr, dimensions):
   operand_avals, init_val_avals = split_list(avals, [len(avals) // 2])
   operand_dtypes = [dtypes.canonicalize_dtype(op.dtype) for op in operand_avals]
@@ -5093,7 +5098,7 @@ def _reduce_jvp_rule(primals, tangents, *, computation, jaxpr, dimensions):
 reduce_p.def_impl(partial(dispatch.apply_primitive, reduce_p))
 reduce_p.def_abstract_eval(
     partial(standard_multi_result_abstract_eval, reduce_p, _reduce_shape_rule,
-            _reduce_dtype_rule, _reduce_weak_type_rule))
+            _reduce_dtype_rule, _reduce_weak_type_rule, _reduce_sharding_rule))
 batching.primitive_batchers[reduce_p] = _reduce_batch_rule
 ad.primitive_jvps[reduce_p] = _reduce_jvp_rule
 
@@ -5115,6 +5120,9 @@ def _reduce_lower(ctx, *values, computation, jaxpr, dimensions):
                                       *reducer.arguments,
                                       dim_var_values=ctx.dim_var_values)
     hlo.return_(mlir.flatten_ir_values(out_nodes))
+  if config.sharding_in_types.value:
+    return [mlir.lower_sharding_under_shit(ctx, r, aval)
+            for r, aval in safe_zip(op.results, ctx.avals_out)]
   return op.results
 
 mlir.register_lowering(reduce_p, _reduce_lower)
@@ -5227,7 +5235,12 @@ def _argminmax_shape_rule(operand, *, axes, index_dtype):
   if operand.shape[axis] < 1:
     raise ValueError("argmin and argmax require non-empty reduced dimension. "
                      f"operand.shape={operand.shape} {axis=}")
-  return tuple(np.delete(operand.shape, axis))
+  return util.tuple_delete(operand.shape, axis)
+
+def _argminmax_sharding_rule(operand, *, axes, index_dtype):
+  axis, = axes
+  return operand.sharding.with_spec(
+      util.tuple_delete(operand.sharding.spec, axis))
 
 def _argminmax_dtype_rule(operand, *, axes, index_dtype):
   if not dtypes.issubdtype(index_dtype, np.integer):
@@ -5264,30 +5277,34 @@ def _compute_argminmax(value_comparator, get_identity,
   # value_comparator is either lax.lt (for argmin) or lax.gt
   # get_identity(operand.dtype) is inf for argmin or -inf for argmax
   axis, = axes
-  indices = broadcasted_iota(index_dtype, np.shape(operand), axis)
+  indices = broadcasted_iota(
+      index_dtype, np.shape(operand), axis,
+      _sharding=operand.sharding if config.sharding_in_types.value else None)
   res = reduce([operand, indices],
                [get_identity(operand.dtype), np.array(0, index_dtype)],
                _ArgMinMaxReducer(value_comparator),
                axes)
   return res[1]
 
 argmin_p = standard_primitive(_argminmax_shape_rule, _argminmax_dtype_rule,
-                              'argmin', weak_type_rule=_strip_weak_type)
+                              'argmin', weak_type_rule=_strip_weak_type,
+                              sharding_rule=_argminmax_sharding_rule)
 batching.defreducer(argmin_p, _get_min_identity)
 ad.defjvp_zero(argmin_p)
 
 argmax_p = standard_primitive(_argminmax_shape_rule, _argminmax_dtype_rule,
-                              'argmax', weak_type_rule=_strip_weak_type)
+                              'argmax', weak_type_rule=_strip_weak_type,
+                              sharding_rule=_argminmax_sharding_rule)
 batching.defreducer(argmax_p, _get_max_identity)
 ad.defjvp_zero(argmax_p)
 
-mlir.register_lowering(argmin_p, mlir.cache_lowering(mlir.lower_fun(
-  partial(_compute_argminmax, lt, _get_min_identity),
-  multiple_results=False)))
+mlir.register_lowering(argmin_p, mlir.cache_lowering(
+    mlir.lower_fun(partial(_compute_argminmax, lt, _get_min_identity),
+                   multiple_results=False)))
 
-mlir.register_lowering(argmax_p, mlir.cache_lowering(mlir.lower_fun(
-  partial(_compute_argminmax, gt, _get_max_identity),
-  multiple_results=False)))
+mlir.register_lowering(argmax_p, mlir.cache_lowering(
+    mlir.lower_fun(partial(_compute_argminmax, gt, _get_max_identity),
+                   multiple_results=False)))
 
 
 def _reduce_logical_shape_rule(operand, *, axes):
@@ -5882,7 +5899,7 @@ def _rng_bit_generator_lowering(
 rng_bit_generator_p.def_abstract_eval(
     partial(standard_multi_result_abstract_eval, rng_bit_generator_p,
             _rng_bit_generator_shape_rule, _rng_bit_generator_dtype_rule,
-            _rng_bit_generator_weak_type_rule))
+            _rng_bit_generator_weak_type_rule, None))
 mlir.register_lowering(rng_bit_generator_p,
                        _rng_bit_generator_lowering)
 

jax/_src/lax/utils.py

@@ -69,16 +69,21 @@ def standard_abstract_eval(prim, shape_rule, dtype_rule, weak_type_rule,
     raise TypeError(avals, least_specialized)
 
 def standard_multi_result_abstract_eval(
-    prim, shape_rule, dtype_rule, weak_type_rule, *avals, **kwargs):
+    prim, shape_rule, dtype_rule, weak_type_rule, sharding_rule,
+    *avals, **kwargs):
   assert prim.multiple_results
   assert all(isinstance(aval, core.UnshapedArray) for aval in avals), avals
   least_specialized = max(map(type, avals), key=_get_array_abstraction_level)
   weak_types = weak_type_rule(*avals, **kwargs)
   if least_specialized is core.ShapedArray:
     out_shapes = shape_rule(*avals, **kwargs)
     out_dtypes = dtype_rule(*avals, **kwargs)
-    return [core.ShapedArray(s, d, weak_type=weak_type)
-            for s, d, weak_type in zip(out_shapes, out_dtypes, weak_types)]
+    out_shardings = (sharding_rule(*avals, **kwargs)
+                     if config.sharding_in_types.value else
+                     [None] * len(out_shapes))
+    return [core.ShapedArray(s, d, weak_type=weak_type, sharding=sh)
+            for s, d, weak_type, sh in zip(out_shapes, out_dtypes, weak_types,
+                                           out_shardings)]
   elif least_specialized is core.UnshapedArray:
     out_dtypes = dtype_rule(*avals, **kwargs)
     return [core.UnshapedArray(dtype, weak_type=weak_type)

tests/pjit_test.py

@@ -5466,6 +5466,28 @@ def g(carry, arr):
         ValueError, "0th dimension of all xs should be replicated"):
       f(carry, jax.device_put(arr, NamedSharding(mesh, P('x', None, None))))
 
+  def test_argminmax(self):
+    mesh = jtu.create_mesh((2, 2), ('x', 'y'))
+    np_inp = np.arange(16.).reshape(8, 2)
+    s = NamedSharding(mesh, P('x', 'y'))
+    arr = jax.device_put(np_inp, s)
+
+    @jax.jit
+    def f(x):
+      z = jnp.argmax(x, axis=0)
+      self.assertEqual(z.sharding.spec, P('y'))
+      a = jnp.argmin(x, axis=1)
+      self.assertEqual(a.sharding.spec, P('x'))
+      return z, a
+
+    out1, out2 = f(arr)
+    self.assertArraysEqual(out1, np.argmax(np_inp, axis=0))
+    self.assertEqual(out1.sharding, NamedSharding(mesh, P('y')))
+    self.assertArraysEqual(out2, np.argmin(np_inp, axis=1))
+    self.assertEqual(out2.sharding, NamedSharding(mesh, P('x')))
+
+    self.assertIn('@Sharding', f.lower(arr).as_text())
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):