Repository: google/jax
Issue: Nested jit/pmap bug that only appears with multiple devices

To reproduce:
```python
import jax
import jax.numpy as np
from jax import lax
from functools import partial

from jax.lib import xla_bridge

device_count = xla_bridge.device_count()

# patch https://github.com/google/jax/pull/1063
from jax.interpreters.pxla import ShardedDeviceTuple
xla_bridge.register_constant_handler(ShardedDeviceTuple, jax.xla._device_tuple_constant_handler)


def multi_step_pmap(state, count):
  @partial(jax.pmap, axis_name='x')
  @jax.jit
  def exchange_and_multi_step(state):
    return state

  @jax.jit
  def time_evolution(state):
    return lax.fori_loop(0, count, lambda i, s: exchange_and_multi_step(s), state)

  return time_evolution(state)

multi_step_pmap(np.zeros((device_count,)), count=1)
```
Here's the traceback:
```
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-4-9a3737aed707> in <module>()
----> 1 multi_step_pmap(np.zeros((8,)), count=1)

13 frames
<ipython-input-3-8106aee94f86> in multi_step_pmap(state, count)
     23     return lax.fori_loop(0, count, lambda i, s: exchange_and_multi_step(s), state)
     24 
---> 25   return time_evolution(state)

google3/third_party/py/jax/api.py in f_jitted(*args, **kwargs)
    129     _check_args(args_flat)
    130     flat_fun, out_tree = flatten_fun_leafout(f, in_tree)
--> 131     out = xla.xla_call(flat_fun, *args_flat, device_values=device_values)
    132     return out if out_tree() is leaf else tree_unflatten(out_tree(), out)
    133 

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **params)
    661   if top_trace is None:
    662     with new_sublevel():
--> 663       ans = primitive.impl(f, *args, **params)
    664   else:
    665     tracers = map(top_trace.full_raise, args)

google3/third_party/py/jax/interpreters/xla.py in _xla_call_impl(fun, *args, **params)
    672 def _xla_call_impl(fun, *args, **params):
    673   device_values = FLAGS.jax_device_values and params.pop('device_values')
--> 674   compiled_fun = _xla_callable(fun, device_values, *map(abstractify, args))
    675   try:
    676     return compiled_fun(*args)

google3/third_party/py/jax/linear_util.py in memoized_fun(f, *args)
    203 
    204   def memoized_fun(f, *args):
--> 205     ans, f_prev = memoized_fun_body(f, args)
    206     if id(f_prev) != id(f):
    207       f.populate_stores(f_prev)

google3/third_party/py/jax/linear_util.py in memoized_fun_body(f, args)
    200   @fastcache.clru_cache(maxsize=max_size)
    201   def memoized_fun_body(f, args):
--> 202     return call(f, *args), f
    203 
    204   def memoized_fun(f, *args):

google3/third_party/py/jax/interpreters/xla.py in _xla_callable(fun, device_values, *abstract_args)
    687     assert not env  # no subtraces here (though cond might eventually need them)
    688     axis_env = AxisEnv(jaxpr_replicas(jaxpr), [], [])
--> 689     compiled, result_shape = _compile_jaxpr(jaxpr, axis_env, consts, *abstract_args)
    690     del master, consts, jaxpr, env
    691   if device_values:

google3/third_party/py/jax/interpreters/xla.py in _compile_jaxpr(jaxpr, axis_env, const_vals, *abstract_args)
    223     raise ValueErrr(msg.format(axis_env.nreps, xb.device_count()))
    224   arg_shapes = list(map(xla_shape, abstract_args))
--> 225   built_c = _jaxpr_computation(jaxpr, axis_env, const_vals, (), *arg_shapes)
    226   result_shape = xla_shape_to_result_shape(built_c.GetReturnValueShape())
    227   return built_c.Compile(arg_shapes, xb.get_compile_options(axis_env.nreps),

google3/third_party/py/jax/interpreters/xla.py in _jaxpr_computation(jaxpr, axis_env, const_vals, freevar_shapes, *arg_shapes)
    291     elif eqn.primitive in initial_style_translations:
    292       rule = initial_style_translations[eqn.primitive]
--> 293       ans = rule(c, axis_env, *in_nodes, **eqn.params)
    294     elif eqn.primitive in parallel_translations:
    295       replica_groups = axis_groups(axis_env, eqn.params['axis_name'])

google3/third_party/py/jax/lax/lax_control_flow.py in _while_loop_translation_rule(c, axis_env, init_val, cond_consts, body_consts, aval_out, cond_jaxpr, body_jaxpr)
    202 
    203   cond_c = xla._jaxpr_computation(cond_jaxpr_converted, axis_env, (), (), shape)
--> 204   body_c = xla._jaxpr_computation(body_jaxpr_converted, axis_env, (), (), shape)
    205   full_ans = c.While(cond_c, body_c, loop_carry)
    206   return c.GetTupleElement(full_ans, 0)

google3/third_party/py/jax/interpreters/xla.py in _jaxpr_computation(jaxpr, axis_env, const_vals, freevar_shapes, *arg_shapes)
    301       env_nodes = list(map(read, const_bindings + freevar_bindings))
    302       rule = call_translations[eqn.primitive]
--> 303       ans = rule(c, subjaxpr, axis_env, env_nodes, in_nodes, **eqn.params)
    304     else:
    305       msg = "XLA translation rule for primitive '{}' not found"

google3/third_party/py/jax/interpreters/pxla.py in _xla_pmap_translation_rule(c, jaxpr, axis_env, env_nodes, in_nodes, axis_name, axis_size)
    591                                 *map(c.GetShape, in_nodes_sharded))
    592   sharded_result = c.Call(subc, env_nodes + in_nodes_sharded)
--> 593   return xla_unshard(c, xla.axis_groups(new_env, axis_name), sharded_result)
    594 xla.call_translations[xla_pmap_p] = _xla_pmap_translation_rule
    595 ad.primitive_transposes[xla_pmap_p] = partial(ad.map_transpose, xla_pmap_p)

google3/third_party/py/jax/interpreters/xla.py in axis_groups(axis_env, name)
    338   else:
    339     mesh_axes = (axis_read(axis_env, name),)
--> 340   return _axis_groups(axis_env.nreps, axis_env.sizes, mesh_axes)
    341 
    342 def _axis_groups(nrep, mesh_spec, mesh_axes):

google3/third_party/py/jax/interpreters/xla.py in _axis_groups(nrep, mesh_spec, mesh_axes)
    342 def _axis_groups(nrep, mesh_spec, mesh_axes):
    343   trailing_size, ragged = divmod(nrep, prod(mesh_spec))
--> 344   assert not ragged
    345   full_spec = list(mesh_spec) + [trailing_size]
    346   iota = onp.arange(prod(full_spec)).reshape(full_spec)

AssertionError: 
```

I did a little debugging -- the `axis_env` object here looks like `AxisEnv(nreps=1, names=['x'], sizes=[8])`

If I remove the outer `jit`, the error goes away.

Changed files:
jax/interpreters/xla.py, tests/pmap_test.py

diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
index 90345f698..6df06f773 100644
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -348,12 +348,19 @@ def _axis_groups(nrep, mesh_spec, mesh_axes):
   return tuple(map(tuple, groups.T))
 
 def jaxpr_replicas(jaxpr):
-  nums = (eqn_replicas(eqn) for eqn in jaxpr.eqns if eqn.bound_subjaxprs)
-  return max(it.chain([1], nums))  # max(itr, default=1)
+  return max(it.chain([1], (eqn_replicas(eqn) for eqn in jaxpr.eqns)))
 
 def eqn_replicas(eqn):
-  (subjaxpr, _, _), = eqn.bound_subjaxprs
-  return eqn.params.get('axis_size', 1) * jaxpr_replicas(subjaxpr)
+  if eqn.bound_subjaxprs:
+    (subjaxpr, _, _), = eqn.bound_subjaxprs
+    return eqn.params.get('axis_size', 1) * jaxpr_replicas(subjaxpr)
+  elif eqn.primitive in initial_style_translations:
+    nums = (jaxpr_replicas(param if type(param) is core.Jaxpr else param.jaxpr)
+            for param in eqn.params.values()
+            if type(param) in (core.Jaxpr, core.TypedJaxpr))
+    return max(it.chain([1], nums))
+  else:
+    return 1
 
 
 def lower_fun(fun, instantiate=False, initial_style=False):
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
index 091a7e0d7..4670e3b56 100644
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -659,6 +659,24 @@ class PmapTest(jtu.JaxTestCase):
 
     f(onp.arange(1.).reshape((1, 1)))  # doesn't crash
 
+  def testIssue1065(self):
+    # from https://github.com/google/jax/issues/1065
+    device_count = xla_bridge.device_count()
+
+    def multi_step_pmap(state, count):
+      @partial(pmap, axis_name='x')
+      @jit
+      def exchange_and_multi_step(state):
+        return state
+
+      @jit
+      def time_evolution(state):
+        return lax.fori_loop(0, count, lambda i, s: exchange_and_multi_step(s), state)
+
+      return time_evolution(state)
+
+    multi_step_pmap(np.zeros((device_count,)), count=1)
+
 
 if __name__ == '__main__':
   absltest.main()