Repository: google/jax

Issue:
BUG: Unbound axis error with xmap + scan + cond.
### Description

I've been getting the error `NameError: unbound axis name: b. The following axis names (e.g. defined by pmap) are available to collective operations: ['x', 'x']` when running the following code (a simplified version of my full code):

```python
import jax
import jax.numpy as jnp
import numpy as onp
from jax.experimental.maps import xmap, Mesh
import tensorflow_probability.substrates.jax as tfp
tfd = tfp.distributions

def loss(key):

  init_z = tfd.Normal(0., 1.).sample(seed=key)

  def scan_fn(prev_z, t):
    new_z = jax.lax.cond(t == 0,
            lambda _: prev_z,
            lambda _: prev_z,
            None)
    return new_z, None

  out, _ = jax.lax.scan(scan_fn, init_z, jnp.arange(10))

  return 0.

def step(key):
  x = loss(key)
  return jnp.mean(x, axis=('b'))

xm_step = xmap(step, in_axes=['b',...], out_axes=[...], axis_resources={'b':'x'})

devices = onp.array(jax.local_devices())

key = jax.random.PRNGKey(0)
keys = jax.random.split(key, num=jax.local_device_count())
with Mesh(devices, ('x',)):
  xm_step(keys)
```

This is strange to me for a few reasons:

1. I am actually binding `b` to `x` using `axis_resources`, I think.
2. `x` appears twice in the available axis names.
3. All of the code in `loss` appears to be necessary to trigger the error. If I remove the sampling of `init_z` or set it to a constant, the error disappears. Similarly if I remove the `scan`, the error disappears. Similarly, if I remove the `cond`, the error disappears.
4. Note that `loss` is returning 0 -- the outcome doesn't depend on the computation in the body.

Any thoughts on why this might occur?

@sharadmv

Full error:

```
Traceback (most recent call last):
  File "/Users/dlaw/dev/jax_bug/run.py", line 34, in <module>
    xm_step(keys)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py", line 1128, in while_loop
    outs = while_p.bind(*cond_consts, *body_consts, *init_vals,
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2339, in bind
    axis_main = max((axis_frame(a).main_trace for a in used_axis_names(self, params)),
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2250, in used_axis_names
    subst_axis_names(primitive, params, subst)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2269, in subst_axis_names
    new_params[name] = subst_axis_names_jaxpr(jaxpr, shadowed_subst)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2325, in subst_axis_names_jaxpr
    subst.axis_names |= used_axis_names_jaxpr(jaxpr)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/_src/util.py", line 273, in wrapped
    result = call(weak_arg, *args, **kwargs)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2320, in used_axis_names_jaxpr
    do_subst_axis_names_jaxpr(jaxpr, subst)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2310, in do_subst_axis_names_jaxpr
    eqns = [subst_axis_names_eqn(eqn, subst, var_map) for eqn in jaxpr.eqns]
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2310, in <listcomp>
    eqns = [subst_axis_names_eqn(eqn, subst, var_map) for eqn in jaxpr.eqns]
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2295, in subst_axis_names_eqn
    outvars = [subst_axis_names_var(v, subst, var_map) for v in eqn.outvars]
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2295, in <listcomp>
    outvars = [subst_axis_names_var(v, subst, var_map) for v in eqn.outvars]
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2285, in subst_axis_names_var
    named_shape = {name: axis_frame(name).size for name in names}
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2285, in <dictcomp>
    named_shape = {name: axis_frame(name).size for name in names}
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2233, in axis_frame
    raise NameError(
jax._src.traceback_util.UnfilteredStackTrace: NameError: unbound axis name: b. The following axis names (e.g. defined by pmap) are available to collective operations: ['x', 'x']

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/dlaw/dev/jax_bug/run.py", line 27, in <module>
    xm_step = xmap(step, in_axes=['b',...], out_axes=[...], axis_resources={'b':'x'})
  File "/Users/dlaw/dev/jax_bug/run.py", line 24, in step
    x = loss(key)
  File "/Users/dlaw/dev/jax_bug/run.py", line 12, in loss
    def scan_fn(prev_z, t):
jax._src.source_info_util.JaxStackTraceBeforeTransformation: NameError: unbound axis name: b. The following axis names (e.g. defined by pmap) are available to collective operations: ['x', 'x']

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/dlaw/dev/jax_bug/run.py", line 34, in <module>
    xm_step(keys)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/experimental/maps.py", line 625, in fun_mapped
    out_flat = xmap_p.bind(fun_flat, *args_flat, **params)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/experimental/maps.py", line 849, in bind
    return core.map_bind(self, fun, *args, in_axes=in_axes, **params)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 2072, in map_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/experimental/maps.py", line 852, in process
    return trace.process_xmap(self, fun, tracers, params)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/core.py", line 687, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/experimental/maps.py", line 653, in xmap_impl
    xmap_callable = make_xmap_callable(
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/linear_util.py", line 295, in memoized_fun
    ans = call(fun, *args)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/experimental/maps.py", line 723, in make_xmap_callable
    return pxla.lower_mesh_computation(
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/_src/profiler.py", line 294, in wrapper
    return func(*args, **kwargs)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/interpreters/pxla.py", line 2559, in lower_mesh_computation
    lowering_result = mlir.lower_jaxpr_to_module(
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/interpreters/mlir.py", line 618, in lower_jaxpr_to_module
    lower_jaxpr_to_fun(
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/interpreters/mlir.py", line 881, in lower_jaxpr_to_fun
    out_vals, tokens_out = jaxpr_subcomp(ctx.replace(name_stack=callee_name_stack),
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/interpreters/mlir.py", line 1008, in jaxpr_subcomp
    ans = rule(rule_ctx, *map(_unwrap_singleton_ir_values, in_nodes),
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/interpreters/mlir.py", line 1071, in f_lowered
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, ctx.avals_in)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/_src/profiler.py", line 294, in wrapper
    return func(*args, **kwargs)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 2092, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 2109, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers_)
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/linear_util.py", line 168, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py", line 362, in _scan_impl
    return _scan_impl_loop(
  File "/Users/dlaw/dev/jax_bug/.env/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py", line 323, in _scan_impl_loop
    _, *outs = while_loop(cond_fun, body_fun, init_val)
NameError: unbound axis name: b. The following axis names (e.g. defined by pmap) are available to collective operations: ['x', 'x']
```


### What jax/jaxlib version are you using?

jax head (0.3.17, 4cb31a6), jaxlib 0.3.15

### Which accelerator(s) are you using?

CPU

### Additional System Info

Python 3.10, Mac OS Monterey, Intel processors


Discussion:
I shrank the repro a bit but haven't yet found the underlying issue:

```python
import jax
import jax.numpy as jnp
import numpy as onp
from jax.experimental.maps import xmap, Mesh
from jax.config import config

config.update('jax_traceback_filtering', 'off')

def loss(init_z):

  def scan_fn(prev_z, t):
    new_z = jax.lax.cond(True,
            lambda _: prev_z,
            lambda _: prev_z,
            None)
    return new_z, None
  out, _ = jax.lax.scan(scan_fn, init_z, jnp.arange(10))
  return out

xm_step = xmap(loss, in_axes=['b',...], out_axes=['b', ...], axis_resources={'b':'x'})

devices = onp.array(jax.local_devices())

with Mesh(devices, ('x',)):
  xm_step(jnp.arange(jax.local_device_count(), dtype=jnp.float32))
```

Resolving patch:
diff --git a/jax/_src/lax/control_flow/conditionals.py b/jax/_src/lax/control_flow/conditionals.py
--- a/jax/_src/lax/control_flow/conditionals.py
+++ b/jax/_src/lax/control_flow/conditionals.py
@@ -671,6 +671,12 @@ def _cond_transpose(reduce_axes, cts, *args, branches, linear):
   assert next(out_iter, None) is None
   return [None] + out
 
+def _cond_axis_substitution(params, subst, traverse):
+  if not traverse:
+    return params
+  branches = tuple(core.subst_axis_names_jaxpr(jaxpr, subst) for jaxpr in params['branches'])
+  return dict(params, branches=branches)
+
 def _cond_typecheck(*in_atoms, branches, linear):
   avals = [x.aval for x in in_atoms]
   tc = partial(_typecheck_param, 'cond')
@@ -753,6 +759,7 @@ def cond_bind(*args, branches, linear):
 batching.axis_primitive_batchers[cond_p] = _cond_batching_rule
 xla.register_initial_style_primitive(cond_p)
 core.custom_typechecks[cond_p] = _cond_typecheck
+core.axis_substitution_rules[cond_p] = _cond_axis_substitution
 pe.partial_eval_jaxpr_custom_rules[cond_p] = _cond_partial_eval_custom
 pe.dce_rules[cond_p] = _cond_dce_rule
 


Test patch: