# Buggy code
def _while_loop_translation_rule(
    c,
    axis_env,
    name_stack,
    avals,
    backend,
    *args,
    cond_jaxpr,
    body_jaxpr,
    cond_nconsts,
    body_nconsts,
):
    cond_consts, body_consts, init_vals = split_list(args, [cond_nconsts, body_nconsts])
    batched = bool(cond_jaxpr.out_avals[0].shape)

    # Since jaxprs don't have tuples and have multiple return values, but we need
    # the HLO While loop to take a single tuple input and output a single boolean
    # (for the cond computation) or a single tuple output (for the body
    # computation), we build XLA computations that handle the tuple munging before
    # generating a Call into the computations formed from the jaxprs.

    init_carry = xops.Tuple(c, cond_consts + body_consts + init_vals)

    cond_c = xb.make_computation_builder("cond_computation")
    cond_carry = xb.parameter(cond_c, 0, c.get_shape(init_carry))
    cond_carry_elts = [xops.GetTupleElement(cond_carry, i) for i in range(len(args))]
    x, _, z = split_list(cond_carry_elts, [cond_nconsts, body_nconsts])
    (pred,) = xla.jaxpr_subcomp(
        cond_c,
        cond_jaxpr.jaxpr,
        backend,
        axis_env,
        _map(partial(xb.constant, cond_c), cond_jaxpr.literals),
        extend_name_stack(name_stack, "cond"),
        *(x + z),
    )
    if batched:
        scalar = ShapedArray((), onp.bool_)
        or_ = xla.primitive_subcomputation(lax.or_p, scalar, scalar)
        pred = xops.Reduce(
            cond_c,
            [pred],
            [xb.constant(cond_c, onp.array(False))],
            or_,
            list(range(cond_jaxpr.out_avals[0].ndim)),
        )

    body_c = xb.make_computation_builder("body_computation")
    body_carry = xb.parameter(body_c, 0, c.get_shape(init_carry))
    body_carry_elts = [xops.GetTupleElement(body_carry, i) for i in range(len(args))]
    x, y, z = split_list(body_carry_elts, [cond_nconsts, body_nconsts])
    new_z = xla.jaxpr_subcomp(
        body_c,
        body_jaxpr.jaxpr,
        backend,
        axis_env,
        _map(partial(xb.constant, body_c), body_jaxpr.literals),
        extend_name_stack(name_stack, "body"),
        *(y + z),
    )
    if batched:
        (body_pred,) = xla.jaxpr_subcomp(
            body_c,
            cond_jaxpr.jaxpr,
            backend,
            axis_env,
            _map(partial(xb.constant, body_c), cond_jaxpr.literals),
            extend_name_stack(name_stack, "body_pred"),
            *(x + z),
        )
        new_z = _map(partial(_pred_bcast_select, body_c, body_pred), new_z, z)
        assert _map(body_c.get_shape, new_z) == _map(
            body_c.get_shape, z
        )  # no broadcast
    new_carry = xops.Tuple(body_c, list(itertools.chain(x, y, new_z)))

    ans = xops.While(cond_c.build(pred), body_c.build(new_carry), init_carry)
    ans_elts = [xops.GetTupleElement(ans, i) for i in range(len(args))]
    _, _, z = split_list(ans_elts, [cond_nconsts, body_nconsts])
    return xops.Tuple(c, z)


# Traceback (AssertionError)
Traceback (most recent call last):
File "test.py", line 21, in <module>
print(vmap_test(arr, arr))
File "/home/adabbott/Git/jax/jax/jax/api.py", line 858, in batched_fun
lambda: flatten_axes(out_tree(), out_axes))
File "/home/adabbott/Git/jax/jax/jax/interpreters/batching.py", line 34, in batch
return batched_fun.call_wrapped(*in_vals)
File "/home/adabbott/Git/jax/jax/jax/linear_util.py", line 150, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "test.py", line 12, in test
for _ in s.while_range(lambda: s.j < b + 1):
File "/home/adabbott/Git/jax/jax/jax/experimental/loops.py", line 341, in __next__
self.end_tracing_body()
File "/home/adabbott/Git/jax/jax/jax/experimental/loops.py", line 407, in end_tracing_body
carried_init_vals, body_typed_jaxpr, body_const_vals)
File "/home/adabbott/Git/jax/jax/jax/experimental/loops.py", line 576, in build_output_vals
body_jaxpr=body_typed_jaxpr)
File "/home/adabbott/Git/jax/jax/jax/core.py", line 212, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "/home/adabbott/Git/jax/jax/jax/interpreters/partial_eval.py", line 141, in process_primitive
return custom_partial_eval_rules[primitive](self, *tracers, **params)
File "/home/adabbott/Git/jax/jax/jax/lax/lax_control_flow.py", line 517, in _while_partial_eval
body_jaxpr=body_jaxpr_known)
File "/home/adabbott/Git/jax/jax/jax/core.py", line 212, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "/home/adabbott/Git/jax/jax/jax/interpreters/batching.py", line 134, in process_primitive
val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
File "/home/adabbott/Git/jax/jax/jax/lax/lax_control_flow.py", line 391, in _while_loop_batching_rule
body_nconsts=body_nconsts, body_jaxpr=body_jaxpr_batched)
File "/home/adabbott/Git/jax/jax/jax/core.py", line 209, in bind
return self.impl(*args, **kwargs)
File "/home/adabbott/Git/jax/jax/jax/interpreters/xla.py", line 217, in apply_primitive
compiled_fun = xla_primitive_callable(prim, *map(arg_spec, args), **params)
File "/home/adabbott/Git/jax/jax/jax/interpreters/xla.py", line 248, in xla_primitive_callable
*avals, **params)
File "/home/adabbott/Git/jax/jax/jax/interpreters/xla.py", line 295, in primitive_computation
*xla_args, **params)
File "/home/adabbott/Git/jax/jax/jax/lax/lax_control_flow.py", line 332, in _while_loop_translation_rule
new_z = _map(partial(_pred_bcast_select, body_c, body_pred), new_z, z)
File "/home/adabbott/Git/jax/jax/jax/util.py", line 34, in safe_map
return list(map(f, *args))
File "/home/adabbott/Git/jax/jax/jax/lax/lax_control_flow.py", line 350, in _pred_bcast_select
assert pred_shape == x_shape[:len(pred_shape)] == y_shape[:len(pred_shape)]
AssertionError

# Fixed code
def _while_loop_translation_rule(
    c,
    axis_env,
    name_stack,
    avals,
    backend,
    *args,
    cond_jaxpr,
    body_jaxpr,
    cond_nconsts,
    body_nconsts,
):
    cond_consts, body_consts, init_vals = split_list(args, [cond_nconsts, body_nconsts])
    batched = bool(cond_jaxpr.out_avals[0].shape)

    # Since jaxprs don't have tuples and have multiple return values, but we need
    # the HLO While loop to take a single tuple input and output a single boolean
    # (for the cond computation) or a single tuple output (for the body
    # computation), we build XLA computations that handle the tuple munging before
    # generating a Call into the computations formed from the jaxprs.

    init_carry = xops.Tuple(c, cond_consts + body_consts + init_vals)

    cond_c = xb.make_computation_builder("cond_computation")
    cond_carry = xb.parameter(cond_c, 0, c.get_shape(init_carry))
    cond_carry_elts = [xops.GetTupleElement(cond_carry, i) for i in range(len(args))]
    x, _, z = split_list(cond_carry_elts, [cond_nconsts, body_nconsts])
    (pred,) = xla.jaxpr_subcomp(
        cond_c,
        cond_jaxpr.jaxpr,
        backend,
        axis_env,
        _map(partial(xb.constant, cond_c), cond_jaxpr.literals),
        extend_name_stack(name_stack, "cond"),
        *(x + z),
    )
    if batched:
        scalar = ShapedArray((), onp.bool_)
        or_ = xla.primitive_subcomputation(lax.or_p, scalar, scalar)
        pred = xops.Reduce(
            cond_c,
            [pred],
            [xb.constant(cond_c, onp.array(False))],
            or_,
            list(range(cond_jaxpr.out_avals[0].ndim)),
        )

    body_c = xb.make_computation_builder("body_computation")
    body_carry = xb.parameter(body_c, 0, c.get_shape(init_carry))
    body_carry_elts = [xops.GetTupleElement(body_carry, i) for i in range(len(args))]
    x, y, z = split_list(body_carry_elts, [cond_nconsts, body_nconsts])
    new_z = xla.jaxpr_subcomp(
        body_c,
        body_jaxpr.jaxpr,
        backend,
        axis_env,
        _map(partial(xb.constant, body_c), body_jaxpr.literals),
        extend_name_stack(name_stack, "body"),
        *(y + z),
    )
    if batched:
        (body_pred,) = xla.jaxpr_subcomp(
            body_c,
            cond_jaxpr.jaxpr,
            backend,
            axis_env,
            _map(partial(xb.constant, body_c), cond_jaxpr.literals),
            extend_name_stack(name_stack, "body_pred"),
            *(x + z),
        )
        new_z = _map(
            partial(_pred_bcast_select, body_c, body_pred),
            new_z,
            z,
            body_jaxpr.out_avals,
        )
        assert _map(body_c.get_shape, new_z) == _map(
            body_c.get_shape, z
        )  # no broadcast
    new_carry = xops.Tuple(body_c, list(itertools.chain(x, y, new_z)))

    ans = xops.While(cond_c.build(pred), body_c.build(new_carry), init_carry)
    ans_elts = [xops.GetTupleElement(ans, i) for i in range(len(args))]
    _, _, z = split_list(ans_elts, [cond_nconsts, body_nconsts])
    return xops.Tuple(c, z)