Repository: google/jax

Issue:
AttributeError: 'Unit' object has no attribute 'shape' while tracing function. What does it mean??
### Discussed in https://github.com/google/jax/discussions/9979

<div type='discussions-op-text'>

<sup>Originally posted by **RadostW** March 21, 2022</sup>
Inside a nested scan when using `jax.lax.cond` I get error message that is hard to understand.
Original code was very complex, sadly minimal example I could make is still quite complex.
Bizzarly, when changing small details code runs fine (see commented lines). 

Working fine:
`x * d_www[0,0,0] # Ok`
`x * f_www(x)[0,0,0,0] # Ok`

Throwing hard to understand error
`jnp.array([1.0]) * f_www(x)[0, 0, 0, 0] * d_www[0, 0, 0]  # AttributeError`

Which makes me think that both `d_www` creation and `f_www` creation work fine.

I've found no similar issues when searching.

~~R


```python
import jax.numpy as jnp
import jax


def upper(y):
    return jnp.exp(y)


def lower(y):
    return jnp.exp(2 * y)


def branch_fun(q):
    x = q[0]
    ret = jax.lax.cond(x > 1.0, upper, lower, x)
    return jnp.array([[ret]])


def L_w(f):
    def wrapped(x):
        return jax.numpy.tensordot(jax.jacobian(f)(x), branch_fun(x), axes=1)

    return wrapped


def solve_many():
    t0 = 0.0
    w0 = jax.numpy.zeros(1)

    id_ = lambda x: x
    f_www = L_w(L_w(L_w(id_)))

    def step(
        x,
        d_www,
    ):

        new_x = (
            # x * d_www[0,0,0] # Ok
            # x * f_www(x)[0,0,0,0] # Ok
            jnp.array([1.0]) * f_www(x)[0, 0, 0, 0] * d_www[0, 0, 0]  # AttributeError
            # x * f_www(x)[0,0,0,0] * d_www[0,0,0] # AttributeError
            # contract_all(f_www(x), d_www) # AttributeError
            # f_www(x)[:,0,0,0] # AttributeError
        )

        return new_x

    key = jax.random.PRNGKey(0)

    def scan_func(x, y):
        xp = step(x, y)
        return (xp, xp)

    def chunk_function(x, y):
        z = jax.lax.scan(scan_func, jnp.array([0.1]), y)[0]
        return z, z

    def get_solution_fragment(starting_state, key):
        wiener_integrals = jax.random.normal(key, shape=(4, 1, 1, 1))

        last_state, solution_values = jax.lax.scan(
            chunk_function,
            starting_state,
            jnp.reshape(wiener_integrals, (2, 2, 1, 1, 1)),
        )

        return (last_state, last_state)

    @jax.vmap
    def get_solution(key):
        _, chunked_solution = jax.lax.scan(
            lambda state, key: get_solution_fragment(state, key),
            (jnp.array([0.1])),
            jax.random.split(key, 1),
        )

        return chunked_solution

    keys = jax.random.split(key, 1)
    solutions = get_solution(keys)

    return solutions


def test_branched_coefficients():
    trajectories = solve_many()


if __name__ == "__main__":
    trajectories = solve_many()
```

```
The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test_branched_coefficients.py", line 91, in <module>
    trajectories = solve_many()
  File "test_branched_coefficients.py", line 81, in solve_many
    solutions = get_solution(keys)
  File "test_branched_coefficients.py", line 72, in get_solution
    _, chunked_solution = jax.lax.scan(
AttributeError: 'Unit' object has no attribute 'shape'
```
</div>


Discussion:
minimal repro:
```python
import jax.numpy as jnp
import jax

id_f = lambda x: x
def b(x):
    return jax.lax.cond(x[0], id_f, id_f, x) # jnp.where(x[0], id_f(x), id_f(x)) is ok

f = jax.jacrev(lambda x: b(b(x))) # jacfwd is ok

def repro():
    jax.vmap(f)(jnp.ones((1,1))) # without vmap is ok

if __name__ == "__main__":
    with jax.disable_jit():
        repro()
```

Resolving patch:
diff --git a/jax/_src/lax/control_flow.py b/jax/_src/lax/control_flow.py
--- a/jax/_src/lax/control_flow.py
+++ b/jax/_src/lax/control_flow.py
@@ -971,10 +971,15 @@ def _cond_batching_rule(axis_size, axis_name, main_type, args, dims, branches, l
     # for the select we broadcast the input operands for simplicity and leave
     # optimizations to XLA.
     # TODO(mattjj,frostig): assumes branches are side-effect-free, revise!
-    index, *ops = [batching.bdim_at_front(x, d, axis_size) for x, d in zip(args, dims)]
+    index, *ops = [
+        batching.bdim_at_front(x, d, axis_size) for x, d in zip(args, dims)]
+
+    in_batched  = [a is not core.abstract_unit for a in branches[0].in_avals]
+    out_batched = [a is not core.abstract_unit for a in branches[0].out_avals]
 
     branches_batched = [
-        batching.batch_jaxpr(jaxpr, axis_size, [True] * len(ops), True, axis_name, main_type)[0]
+        batching.batch_jaxpr(
+            jaxpr, axis_size, in_batched, out_batched, axis_name, main_type)[0]
         for jaxpr in branches]
 
     branch_outs = []
@@ -987,7 +992,7 @@ def _cond_batching_rule(axis_size, axis_name, main_type, args, dims, branches, l
       branch_outs.append(core.jaxpr_as_fun(jaxpr)(*ops_))
     out = [_bcast_select_n(index, *outs) if outs[0] is not core.unit else outs[0]
            for outs in zip(*branch_outs)]
-    return out, [0] * len(branch_outs[0])
+    return out, [0 if b else None for b in out_batched]
   else:
     ops_bat = [d is not batching.not_mapped for d in op_dims]
     ops = [batching.moveaxis(x, d, 0) if b else x
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -118,6 +118,7 @@ def __init__(self, trace, val, batch_dim: Optional[int],
       if type(batch_dim) is int:
         aval = raise_to_shaped(core.get_aval(val))
         assert aval is core.abstract_unit or 0 <= batch_dim < len(aval.shape)  # type: ignore
+        assert aval is not core.abstract_unit or batch_dim is not_mapped
     self._trace = trace
     self.val = val
     self.batch_dim = batch_dim


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -2393,5 +2393,11 @@ def body(carry):
       return lax.while_loop(cond, body, (i, jnp.ones(3)))[1]
     jax.vmap(f, in_axes=(0, 1))(jnp.arange(4), jnp.ones((3, 4)))
 
+  def test_cond_ad_batched_unit(self):
+    # see issue #9985
+    def cond_id(x):
+      return lax.cond(x < 0., lambda x: x, lambda x: x, x)
+    jax.vmap(jax.jacrev(lambda x: cond_id(cond_id(x))))(jnp.ones(1))
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())