# Buggy code
def _pred_bcast_select(c, pred, x, y):
    pred_shape = c.get_shape(pred).dimensions()
    x_shape = c.get_shape(x).dimensions()
    y_shape = c.get_shape(y).dimensions()
    assert x_shape == y_shape
    if not c.get_shape(x).is_array() and not c.get_shape(y).is_array():
        # Two tokens
        return xops.AfterAll(c, [x, y])
    else:
        assert pred_shape == x_shape[: len(pred_shape)] == y_shape[: len(pred_shape)]
        bcast_pred = xops.BroadcastInDim(pred, x_shape, list(range(len(pred_shape))))
        return xops.Select(bcast_pred, x, y)


# 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 _pred_bcast_select(c, pred, x, y, x_y_aval: core.AbstractValue):
    pred_shape = c.get_shape(pred).dimensions()
    x_shape = c.get_shape(x).dimensions()
    y_shape = c.get_shape(y).dimensions()
    assert x_shape == y_shape
    if x_y_aval is core.abstract_unit:
        return x
    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)]
        bcast_pred = xops.BroadcastInDim(pred, x_shape, list(range(len(pred_shape))))
        return xops.Select(bcast_pred, x, y)