Repository: google/jax
Issue: weak_types: error in grad of switch

This is an error on the current master that was introduced by #4161. In a backward pass over a switch statement, the in_avals and out_avals do not match if branches return data of a different weak type.

This is the same error you would see if you had branches that returned different dtypes, but here the dtypes match but the weak_types do not.

Short repro:
```python
from jax import grad, lax
import jax.numpy as jnp

branches = [
    lambda x: x,           # This preserves the weak type of x.
    lambda x: jnp.cos(x),  # This strips the weak type of x.
]

def f(x):
  return lax.switch(x.astype(jnp.int32), branches, x)

grad(f)(0.0)
```
```pytb
Traceback (most recent call last):
  File "tmp.py", line 12, in <module>
    grad(f)(0.0)
  File "/Users/vanderplas/github/google/jax/jax/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/vanderplas/github/google/jax/jax/api.py", line 750, in grad_f
    _, g = value_and_grad_f(*args, **kwargs)
  File "/Users/vanderplas/github/google/jax/jax/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/vanderplas/github/google/jax/jax/api.py", line 819, in value_and_grad_f
    g = vjp_py(np.ones((), dtype=dtype))
  File "/Users/vanderplas/github/google/jax/jax/api.py", line 1791, in _vjp_pullback_wrapper
    ans = fun(*args)
  File "/Users/vanderplas/github/google/jax/jax/interpreters/ad.py", line 120, in unbound_vjp
    arg_cts = backward_pass(jaxpr, consts, dummy_args, cts)
  File "/Users/vanderplas/github/google/jax/jax/interpreters/ad.py", line 220, in backward_pass
    cts_out = get_primitive_transpose(eqn.primitive)(cts_in, *invals,
  File "/Users/vanderplas/github/google/jax/jax/_src/lax/control_flow.py", line 1023, in _cond_transpose
    assert all(jaxpr.out_avals == lin_in_avals for jaxpr in branches_trans)
```

Changed files:
tests/lax_control_flow_test.py, jax/_src/lax/control_flow.py, jax/core.py

diff --git a/jax/_src/lax/control_flow.py b/jax/_src/lax/control_flow.py
index b68ac87b0..ad5464de0 100644
--- a/jax/_src/lax/control_flow.py
+++ b/jax/_src/lax/control_flow.py
@@ -1021,8 +1021,11 @@ def _cond_transpose(cts, *args, branches, linear):
 
   branches_trans = tuple(
       _transpose_cond_jaxpr(jaxpr, num_res) for jaxpr in branches)
-  lin_in_avals = [raise_to_shaped(a, weak_type=False) for a, l in zip(in_avals, linear) if l]
-  assert all(jaxpr.out_avals == lin_in_avals for jaxpr in branches_trans)
+  lin_in_avals = [raise_to_shaped(a, weak_type=False)
+                  for a, l in zip(in_avals, linear) if l]
+  assert all(core.typematch(out_aval, lin_in_aval)
+             for jaxpr in branches_trans
+             for out_aval, lin_in_aval in zip(jaxpr.out_avals, lin_in_avals))
 
   res = ops[:num_res]
   cts = _map(ad.instantiate_zeros_aval, branches[0].out_avals, cts)
diff --git a/jax/core.py b/jax/core.py
index 6699b1b2d..170f5b6c7 100644
--- a/jax/core.py
+++ b/jax/core.py
@@ -1317,7 +1317,8 @@ def check_jaxpr(jaxpr: Jaxpr):
   - variables are typed equally throughout a jaxpr
   - variable type annotations are compatible with their binding expression
 
-  Raises `TypeError` if `jaxpr` is determined invalid. Returns `None` otherwise.
+  Raises `JaxprTypeError` if `jaxpr` is determined invalid. Returns `None`
+  otherwise.
   """
   try:
     _check_jaxpr(jaxpr, [v.aval for v in jaxpr.invars])
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
index f59d52693..274aee925 100644
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -999,6 +999,24 @@ class LaxControlFlowTest(jtu.JaxTestCase):
       self.assertAllClose(ans, expected, check_dtypes=False)
       jtu.check_grads(f, (x,), order=2, modes=["fwd", "rev"])
 
+  def testSwitchGradWithWeakTypeMismatch(self):  # issue 4696
+    branches = [
+        lambda x: x,           # This preserves the weak type of x.
+        lambda x: jnp.cos(x),  # This strips the weak type of x.
+    ]
+
+    def f_ref(x):
+      i = x.astype(jnp.int32)
+      return branches[i](x)
+
+    def f(x):
+      return lax.switch(x.astype(jnp.int32), branches, x)
+
+    for x in [0., 1.]:
+      ans = api.grad(f)(x)
+      expected = api.grad(f_ref)(x)
+      self.assertAllClose(ans, expected, check_dtypes=False)
+
   @parameterized.named_parameters(
       {"testcase_name": f"_{name}", "cond": cond}
       for cond, name in COND_IMPLS)