# Buggy code
def eqn_replicas(eqn):
    ((subjaxpr, _, _),) = eqn.bound_subjaxprs
    return eqn.params.get("axis_size", 1) * jaxpr_replicas(subjaxpr)


# Traceback (AssertionError)
---------------------------------------------------------------------------
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:

# Fixed code
def eqn_replicas(eqn):
    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