[Take 2] Merge all_gather with all_gather_reduced, psum_scatter with unreduced_psum_scatter and psum with unreduced_psum.

Here are the changes:

all_gather signature gets a to argument. all_gather(x, axis_name, tiled=True, to=...). The allowed values are varying and reduced. to defaults to varying to preserve the current behavior but you can get AGR by specifying to='reduced'

psum_scatter will infer the input state from the type. If the input is unreduced over the axis_name, then we will dispatch to unreduced_psum_scatter_p. If the input is varying, it will dispatch to reduce_scatter_p

psum will infer the input state from the type. If the input is unreduced over the axis_name, then we will dispatch to unreduced_psum_p. If the input is varying, it will dispatch to psum_invariant_p

Reverts 5b9cfa3

PiperOrigin-RevId: 839418947

Author: yashk2810

jax/_src/lax/parallel.py

@@ -121,6 +121,21 @@ def psum(x, axis_name, *, axis_index_groups=None):
      [20 22 24 26]
      [20 22 24 26]]
   """
+  axes = ((axis_name,) if not isinstance(axis_name, (tuple, list)) else
+          tuple(axis_name))
+  if not axes:
+    return x
+  def bind(leaf):
+    from_ = _get_from(core.typeof(leaf), axes, 'jax.lax.psum')
+    if from_ == 'unreduced':
+      if axis_index_groups is not None:
+        raise NotImplementedError
+      return unreduced_psum(leaf, axes)
+    else:
+      return _psum(leaf, axes, axis_index_groups=axis_index_groups)
+  return tree_util.tree_map(bind, x)
+
+def _psum(x, axis_name, *, axis_index_groups):
   if not isinstance(axis_name, (tuple, list)):
     axis_name = (axis_name,)
   if not axis_name:
@@ -1611,7 +1626,8 @@ def insert_collective_pvary(axis_name, x):
   x = pvary(x, tuple(n for n in names_union if n not in aval.vma))
   return x
 
-def all_gather(x, axis_name, *, axis_index_groups=None, axis=0, tiled=False):
+def all_gather(x, axis_name, *, axis_index_groups=None, axis=0, tiled=False,
+               to: str = 'varying'):
   """Gather values of x across all replicas.
 
   If ``x`` is a pytree then the result is equivalent to mapping this function to
@@ -1675,6 +1691,22 @@ def all_gather(x, axis_name, *, axis_index_groups=None, axis=0, tiled=False):
    [[12 13 14 15]
     [ 4  5  6  7]]]
   """
+  _allowed_ag_to = {'varying', 'reduced'}
+  if to not in _allowed_ag_to:
+    raise ValueError(
+        "Got unexpected `to` value for `jax.lax.all_gather`. Allowed `to`"
+        f" values are: {_allowed_ag_to}")
+  if to == 'varying':
+    return _all_gather(x, axis_name, axis_index_groups=axis_index_groups,
+                       axis=axis, tiled=tiled)
+  else:
+    assert to == 'reduced'
+    if axis_index_groups is not None:
+      raise NotImplementedError
+    return all_gather_reduced(x, axis_name, axis=axis, tiled=tiled)
+
+
+def _all_gather(x, axis_name, *, axis_index_groups, axis, tiled):
   if not isinstance(axis_name, tuple):
     axis_name = (axis_name,)
   if not axis_name:
@@ -2131,6 +2163,22 @@ def psum_scatter(x, axis_name, *, scatter_dimension=0, axis_index_groups=None,
    [12 14]
    [16 18]]
   """
+  axes = (axis_name,) if not isinstance(axis_name, tuple) else axis_name
+  if not axes:
+    return x
+  def bind(leaf):
+    from_ = _get_from(core.typeof(leaf), axes, 'jax.lax.psum_scatter')
+    if from_ == 'unreduced':
+      if axis_index_groups is not None:
+        raise NotImplementedError
+      return unreduced_psum_scatter(
+          leaf, axes, scatter_dimension=scatter_dimension, tiled=tiled)
+    else:
+      return _psum_scatter(leaf, axes, scatter_dimension=scatter_dimension,
+                           axis_index_groups=axis_index_groups, tiled=tiled)
+  return tree_util.tree_map(bind, x)
+
+def _psum_scatter(x, axis_name, *, scatter_dimension, axis_index_groups, tiled):
   if not isinstance(axis_name, tuple):
     axis_name = (axis_name,)
   if not axis_name:
@@ -2744,7 +2792,7 @@ def _reduced_vary_cast_batcher(vals_in, dims_in, *, axes):
 
 ################################## pcast #############################
 
-def _get_from(aval, axes: tuple[AxisName, ...]) -> str:
+def _get_from(aval, axes: tuple[AxisName, ...], name) -> str:
   vma = aval.vma
   unreduced = aval.sharding.spec.unreduced
   reduced = aval.sharding.spec.reduced
@@ -2765,7 +2813,7 @@ def _get_from(aval, axes: tuple[AxisName, ...]) -> str:
 
   if len(out) > 1:
     raise ValueError(
-        "`jax.lax.pcast` can only accept axis_name which corresponds to one of"
+        f"{name} can only accept axis_name which corresponds to one of"
         " varying, unreduced, reduced or invarying state of the input. Got"
         f" input type: {aval}, axes: {axes} and input state: {out}")
   o, = out
@@ -2779,18 +2827,22 @@ def _get_from(aval, axes: tuple[AxisName, ...]) -> str:
     ('reduced', 'varying'): core.reduced_vary_cast,
 }
 
-_allowed_to = {'unreduced', 'reduced', 'varying'}
+_allowed_pcast_to = {'unreduced', 'reduced', 'varying'}
 
 def pcast(x, axis_name, *, to: str):
   axes = (axis_name,) if not isinstance(axis_name, tuple) else axis_name
   if not axis_name:
     return x
 
-  if to not in _allowed_to:
+  if to not in _allowed_pcast_to:
     raise ValueError(
-        f'Got unexpected `to` value. Allowed `to` values are: {_allowed_to}')
-  from_ = _get_from(core.typeof(x), axes)
-  func = _pcast_funcs.get((from_.lower(), to.lower()), None)
-  if func is None:
-    raise ValueError(f"Unsupported pcast from={from_}, {to=}")
-  return func(x, axes)
+        "Got unexpected `to` value. Allowed `to` values are:"
+        f" {_allowed_pcast_to}")
+
+  def bind(leaf):
+    from_ = _get_from(core.typeof(leaf), axes, 'jax.lax.pcast')
+    func = _pcast_funcs.get((from_, to), None)
+    if func is None:
+      raise ValueError(f"Unsupported pcast from={from_}, {to=}")
+    return func(leaf, axes)
+  return tree_util.tree_map(bind, x)

jax/lax/__init__.py

@@ -357,9 +357,6 @@
 )
 from jax._src.lax.parallel import (
   all_gather as all_gather,
-  all_gather_reduced as all_gather_reduced,
-  unreduced_psum_scatter as unreduced_psum_scatter,
-  unreduced_psum as unreduced_psum,
   pcast as pcast,
   all_gather_p as all_gather_p,
   all_to_all as all_to_all,

tests/shard_map_test.py

@@ -2598,7 +2598,7 @@ def f(x, y):
       def ag(a):
         self.assertEqual(a.aval.vma, {'seq'})
         self.assertEqual(a.aval.sharding.spec.unreduced, frozenset())
-        out = lax.all_gather_reduced(a, axis_name='seq', tiled=True)
+        out = lax.all_gather(a, axis_name='seq', tiled=True, to='reduced')
         self.assertEqual(out.aval.vma, frozenset())
         self.assertEqual(out.aval.sharding.spec.unreduced, frozenset())
         self.assertEqual(out.aval.sharding.spec.reduced, {'seq'})
@@ -2618,7 +2618,7 @@ def f_bwd(res, g):
       def rs(a):
         self.assertEqual(a.aval.vma, frozenset())
         self.assertEqual(a.aval.sharding.spec.unreduced, {'data', 'seq'})
-        out = lax.unreduced_psum_scatter(a, axis_name='seq', tiled=True)
+        out = lax.psum_scatter(a, axis_name='seq', tiled=True)
         self.assertEqual(out.aval.vma, {'seq'})
         self.assertEqual(out.aval.sharding.spec.unreduced, {'data'})
         return out
@@ -2628,7 +2628,7 @@ def rs(a):
       def ar(a):
         self.assertEqual(a.aval.vma, {'seq'})
         self.assertEqual(a.aval.sharding.spec.unreduced, {'data'})
-        out = lax.unreduced_psum(a, axis_name='data')
+        out = lax.psum(a, axis_name='data')
         self.assertEqual(out.aval.vma, {'seq'})
         self.assertEqual(out.aval.sharding.spec.unreduced, frozenset())
         return out
@@ -2684,7 +2684,7 @@ def ag(a):
       self.assertEqual(a.aval.vma, {'seq'})
       self.assertEqual(a.aval.sharding.spec.unreduced, frozenset())
       self.assertEqual(a.aval.sharding.spec.reduced, {'data'})
-      out = lax.all_gather_reduced(a, axis_name='seq', tiled=True)
+      out = lax.all_gather(a, axis_name='seq', tiled=True, to='reduced')
       self.assertEqual(out.aval.vma, frozenset())
       self.assertEqual(out.aval.sharding.spec.unreduced, frozenset())
       self.assertEqual(out.aval.sharding.spec.reduced, {'seq', 'data'})
@@ -2727,7 +2727,7 @@ def test_unreduced_psum_fwd_preduced_bwd(self, mesh):
     def ar(x):
       self.assertEqual(x.aval.vma, frozenset())
       self.assertEqual(x.aval.sharding.spec.unreduced, {'x'})
-      out = jax.lax.unreduced_psum(x, 'x')
+      out = jax.lax.psum(x, 'x')
       self.assertEqual(out.aval.vma, frozenset())
       self.assertEqual(out.aval.sharding.spec.unreduced, frozenset())
       return out
@@ -2797,7 +2797,7 @@ def test_all_gather_reduced_fwd_unreduced_psum_scatter_bwd(self, mesh):
     def ag(a):
       self.assertEqual(a.aval.vma, {'seq'})
       self.assertEqual(a.aval.sharding.spec.unreduced, frozenset())
-      out = lax.all_gather_reduced(a, axis_name='seq', tiled=True)
+      out = lax.all_gather(a, axis_name='seq', tiled=True, to='reduced')
       self.assertEqual(out.aval.vma, frozenset())
       self.assertEqual(out.aval.sharding.spec.unreduced, frozenset())
       self.assertEqual(out.aval.sharding.spec.reduced, {'seq'})
@@ -2834,7 +2834,7 @@ def test_unreduced_psum_scatter_fwd_all_gather_reduced_bwd(self, mesh):
     def rs(a):
       self.assertEqual(a.aval.vma, frozenset())
       self.assertEqual(a.aval.sharding.spec.unreduced, {'x'})
-      out = lax.unreduced_psum_scatter(a, axis_name='x', tiled=True)
+      out = lax.psum_scatter(a, axis_name='x', tiled=True)
       self.assertEqual(out.aval.vma, {'x'})
       self.assertEqual(out.aval.sharding.spec.unreduced, frozenset())
       return out
@@ -4469,7 +4469,7 @@ def f(x, y):
       return jax.lax.pcast(a, ('x', 'y'), to='reduced')
 
     with self.assertRaisesRegex(
-        ValueError, "`jax.lax.pcast` can only accept axis_name which"):
+        ValueError, "jax.lax.pcast can only accept axis_name which"):
       f(arr1, arr2)
 
   @parameterized.named_parameters(