Repository: google/jax

Issue:
Loops don't track axis names correctly
As reported by @sharadmv:
```py
import jax
from jax import lax
from jax.experimental import maps

def f(x):
  z = x + lax.axis_index('a')
  y = x + lax.axis_index('b')
  def cond(carry):
    i, x = carry
    return x < 5
  def body(carry):
    i, x = carry
    return i + 1, x + lax.psum(y, 'b')
  return lax.while_loop(cond, body, (0, z))[1]
maps.xmap(f, axis_sizes=dict(a=2, b=10), out_axes=(['a']), in_axes={})(1.)
```


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/control_flow.py b/jax/_src/lax/control_flow.py
--- a/jax/_src/lax/control_flow.py
+++ b/jax/_src/lax/control_flow.py
@@ -373,7 +373,7 @@ def _pred_bcast_select(c, pred, x, y, x_y_aval: core.AbstractValue):
   elif x_y_aval is core.abstract_token:
     return xops.AfterAll(c, [x, y])
   else:
-    assert pred_shape == x_shape[:len(pred_shape)] == y_shape[:len(pred_shape)]
+    assert pred_shape == x_shape[:len(pred_shape)] == y_shape[:len(pred_shape)], (pred_shape, x_shape, y_shape)
     bcast_pred = xops.BroadcastInDim(pred, x_shape, list(range(len(pred_shape))))
     return xops.Select(bcast_pred, x, y)
 
@@ -383,43 +383,76 @@ def _while_loop_batching_rule(args, dims, axis_name, main_type,
   size, = {x.shape[d] for x, d in zip(args, dims) if d is not batching.not_mapped}
   orig_batched = [d is not batching.not_mapped for d in dims]
   cconst_bat, bconst_bat, init_bat = split_list(orig_batched, [cond_nconsts, body_nconsts])
+  cconsts, bconsts, init = split_list(args, [cond_nconsts, body_nconsts])
+  cconst_dims, bconst_dims, init_dims = split_list(dims, [cond_nconsts, body_nconsts])
 
+  carry_bat = init_bat
   # Fixpoint computation of which carry are batched: either
   # batched from init, or the carry out is batched. Each iteration promotes
-  # at least one carry to batched. We need at most len(carry) iterations,
-  # but we need one last iteration to prepare the jaxpr based on the final
-  # carry_bat.
-  carry_bat = init_bat
+  # at least one carry to batched. We need at most len(carry) iterations to
+  # reach a fixpoint.
   for _ in range(1 + len(carry_bat)):
-    batched = bconst_bat + carry_bat
-    body_jaxpr_batched, carry_bat_out = batching.batch_jaxpr(
-        body_jaxpr, size, batched, instantiate=carry_bat,
+    _, carry_bat_out = batching.batch_jaxpr(
+        body_jaxpr, size, bconst_bat + carry_bat, instantiate=False,
         axis_name=axis_name, main_type=main_type)
-    cond_jaxpr_batched, (pred_bat,) = batching.batch_jaxpr(
-        cond_jaxpr, size, cconst_bat + carry_bat,
-        instantiate=bool(cond_jaxpr.out_avals[0].shape),
-        axis_name=axis_name, main_type=main_type)
-    carry_bat_out = _map(partial(operator.or_, pred_bat), carry_bat_out)
-    if carry_bat_out == carry_bat:
+    if carry_bat == carry_bat_out:
       break
-    else:
-      carry_bat = _map(operator.or_, carry_bat, carry_bat_out)
+    carry_bat = safe_map(operator.or_, carry_bat, carry_bat_out)
   else:
     assert False, "Fixpoint not reached"
 
-  consts, init = split_list(args, [cond_nconsts + body_nconsts])
-  const_dims, init_dims = split_list(dims, [cond_nconsts + body_nconsts])
-  new_consts = [batching.moveaxis(x, d, 0) if d is not batching.not_mapped and d != 0
-                else x for x, d in zip(consts, const_dims)]
-  new_init = [batching.broadcast(x, size, 0) if now_bat and not was_bat
-              else batching.moveaxis(x, d, 0) if now_bat and d != 0 else x
-              for x, d, was_bat, now_bat in zip(init, init_dims, init_bat, carry_bat)]
+  # Knowing how the carry is batched now, we can determine if the predicate is
+  # batched.
+  _, (pred_bat,) = batching.batch_jaxpr(
+      cond_jaxpr, size, cconst_bat + carry_bat, instantiate=False,
+      axis_name=axis_name, main_type=main_type)
+
+  if pred_bat:
+    # If the predicate is batched, we have to batch *all* of the carry
+    # regardless of if the body needs it.
+    carry_bat = [True] * len(carry_bat)
+    carry_dims = [0] * len(carry_bat)
+    body_jaxpr_batched, _ = batching.batch_jaxpr_axes(
+        body_jaxpr, size, bconst_dims + carry_dims,
+        carry_dims, axis_name=axis_name, main_type=main_type)
+    cond_jaxpr_batched, _ = batching.batch_jaxpr_axes(
+        cond_jaxpr, size, cconst_dims + carry_dims, [0],
+        axis_name=axis_name, main_type=main_type)
+  else:
+    # If the predicate is not batched, we can look at the `cond_jaxpr`'s out
+    # shape to determine the rank of the predicate. From this rank
+    # we pick the dims of the carry to be batched to ensure that the predicate
+    # shape is a prefix of the carry in and out shapes. We can then batch
+    # the `body_jaxpr` according to these new batch dims.
+    cond_rank = len(cond_jaxpr.out_avals[0].shape)
+    carry_dims = [cond_rank if b else None for b in carry_bat]
+    body_jaxpr_batched, _ = batching.batch_jaxpr_axes(
+        body_jaxpr, size, bconst_dims + carry_dims, carry_dims,
+        axis_name=axis_name, main_type=main_type)
+    # Now we need to rebatch the `cond_jaxpr` according to the new dims of the
+    # carry.
+    cond_jaxpr_batched, _ = batching.batch_jaxpr_axes(
+        cond_jaxpr, size, cconst_dims + carry_dims, (None,),
+        axis_name=axis_name, main_type=main_type)
+
+  # To prepare the `init` to the `while_p`, we broadcast values if they are
+  # unbatched and need to have an out axis. If their current batch axis does not
+  # match the one it needs to be for the translation rule to work, we move it
+  # into place.
+  new_init = []
+  for x, old_axis, new_axis in zip(init, init_dims, carry_dims):
+    if old_axis is batching.not_mapped and new_axis is not batching.not_mapped:
+      new_init.append(batching.broadcast(x, size, new_axis))
+    elif old_axis is batching.not_mapped and new_axis is batching.not_mapped:
+      new_init.append(x)
+    else:
+      assert new_axis is not batching.not_mapped
+      new_init.append(batching.moveaxis(x, old_axis, new_axis))
 
-  outs = while_p.bind(*(new_consts + new_init),
+  outs = while_p.bind(*(cconsts + bconsts + new_init),
                       cond_nconsts=cond_nconsts, cond_jaxpr=cond_jaxpr_batched,
                       body_nconsts=body_nconsts, body_jaxpr=body_jaxpr_batched)
-  out_bdims = [0 if b else batching.not_mapped for b in carry_bat]
-  return outs, out_bdims
+  return outs, carry_dims
 
 def _while_loop_jvp(primals, tangents, cond_nconsts, cond_jaxpr, body_nconsts,
                     body_jaxpr):
@@ -551,7 +584,7 @@ def _while_transpose_error(*_, **kwargs):
                    "lax.while_loop or lax.fori_loop. "
                    "Try using lax.scan instead.")
 
-while_p = lax.Primitive('while')
+while_p = core.Primitive('while')
 while_p.multiple_results = True
 while_p.def_impl(partial(xla.apply_primitive, while_p))
 while_p.def_abstract_eval(_while_loop_abstract_eval)
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -470,17 +470,31 @@ def bdim_at_front(x, bdim, size):
     return moveaxis(x, bdim, 0)
 
 
+zero_if_mapped = object()
+
 def batch_jaxpr(closed_jaxpr, axis_size, in_batched, instantiate, axis_name, main_type):
+  if instantiate is None:
+    instantiate = False
+  if isinstance(instantiate, bool):
+    instantiate = [instantiate] * len(closed_jaxpr.out_avals)
+  out_axes = [0 if inst else zero_if_mapped for inst in instantiate]
+  return batch_jaxpr_axes(
+      closed_jaxpr, axis_size,
+      [0 if b else not_mapped for b in in_batched],
+      out_axes,
+      axis_name, main_type)
+
+def batch_jaxpr_axes(closed_jaxpr, axis_size, in_axes, out_axes, axis_name, main_type):
   f = lu.wrap_init(core.jaxpr_as_fun(closed_jaxpr))
-  f, out_batched = batch_subtrace_instantiate(f, instantiate, axis_size)
-  f = batchfun(f, axis_name, axis_size, [0 if b else None for b in in_batched], main_type)
-  avals_in = [core.unmapped_aval(axis_size, 0, aval) if b else aval
-              for aval, b in zip(closed_jaxpr.in_avals, in_batched)]
+  f, out_batched = batch_subtrace_instantiate(f, axis_size, out_axes)
+  f = batchfun(f, axis_name, axis_size, in_axes, main_type)
+  avals_in = [core.unmapped_aval(axis_size, b, aval) if b is not not_mapped
+              else aval for aval, b in zip(closed_jaxpr.in_avals, in_axes)]
   jaxpr_out, _, consts = pe.trace_to_jaxpr_dynamic(f, avals_in)
   return core.ClosedJaxpr(jaxpr_out, consts), out_batched()
 
 @lu.transformation_with_aux
-def batch_subtrace_instantiate(instantiate, axis_size, main, in_dims, *in_vals):
+def batch_subtrace_instantiate(axis_size, out_axes, main, in_dims, *in_vals):
   # this is like `batch_subtrace` but we take an extra `instantiate` arg
   # analogue of `jvp_subtrace` in ad.py
   trace = main.with_cur_sublevel()
@@ -490,13 +504,12 @@ def batch_subtrace_instantiate(instantiate, axis_size, main, in_dims, *in_vals):
   out_tracers = map(trace.full_raise, outs)
   out_vals, out_dims = unzip2((t.val, t.batch_dim) for t in out_tracers)
 
-  if type(instantiate) is bool:
-    instantiate = [instantiate] * len(out_vals)
-  out_vals = [moveaxis(x, d, 0) if d is not not_mapped and d != 0
-              else broadcast(x, axis_size, 0) if d is not_mapped and inst else x
-              for x, d, inst in zip(out_vals, out_dims, instantiate)]
-  out_batched = [d is not not_mapped or inst
-                 for d, inst in zip(out_dims, instantiate)]
+  out_axes = [(None if od is not_mapped else 0) if out_axis is zero_if_mapped else out_axis
+              for od, out_axis in zip(out_dims, out_axes)]
+  out_vals = [moveaxis(x, d, od) if d is not not_mapped
+              else broadcast(x, axis_size, od) if od is not None else x
+              for x, d, od in zip(out_vals, out_dims, out_axes)]
+  out_batched = [od is not None for od in out_axes]
   yield out_vals, out_batched
 
 @lu.transformation_with_aux


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -36,6 +36,7 @@
 from jax import test_util as jtu
 from jax import tree_util
 from jax._src.util import unzip2
+from jax.experimental import maps
 from jax.lib import xla_bridge
 from jax.interpreters import xla
 import jax.numpy as jnp  # scan tests use numpy
@@ -2739,5 +2740,30 @@ def side_effecting_scan(carry, val):
       lax.scan(side_effecting_scan, None, jnp.ones((2, 2)))
       lst[0] += 1
 
+  def test_while_loop_fixed_point_with_nested_named_axes(self):
+    def f(x):
+      z = x + lax.axis_index('a')
+      y = x + lax.axis_index('b')
+      def cond(carry):
+        i, x = carry
+        return x < 5
+      def body(carry):
+        i, x = carry
+        return i + 1, x + lax.psum(y, 'b')
+      return lax.while_loop(cond, body, (0, z))[1]
+    maps.xmap(f, axis_sizes=dict(a=2, b=10), out_axes=(['a']), in_axes={})(1.)
+
+  def test_while_loop_fixed_point_with_batched_pred_and_consts(self):
+    def f(i, x):
+      def cond(carry):
+        i, x = carry
+        return i < 5
+      def body(carry):
+        i, z = carry
+        # Close over const with batch dim = 1
+        return i + 1, z + x
+      return lax.while_loop(cond, body, (i, jnp.ones(3)))[1]
+    jax.vmap(f, in_axes=(0, 1))(jnp.arange(4), jnp.ones((3, 4)))
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())