Repository: google/jax

Issue:
custom_jvp + cond on tangent side = assertion failure
### Description

from @bunelr

```python
import jax
import jax.numpy as jnp

def mask_fun(arr, choice):
  out = (1 - choice) * arr.sum() +  choice * (1 - arr.sum())
  return out

def switch_fun(arr, choice):
  choice = jnp.floor(choice).astype(jnp.int32)
  out = jax.lax.switch(choice, [lambda x: x.sum(), lambda x: 1 - x.sum()], arr)
  return out

test_arr = jnp.arange(3.)
test_val = 0.

print(mask_fun(test_arr, test_val))     # 3.0, all good
print(switch_fun(test_arr, test_val))   # 3.0, all good

print(jax.grad(mask_fun)(test_arr, test_val))   # -> 1. everywhere -> good
print(jax.grad(switch_fun)(test_arr, test_val)) # -> 1. everywhere -> good

def good_switchfun_jvp(primals, tangents):
  arr, choice = primals
  arr_dot, choice_dot = tangents
  return switch_fun(arr, choice), mask_fun(arr_dot, choice)

def bad_switchfun_jvp(primals, tangents):
  arr, choice = primals
  arr_dot, choice_dot = tangents
  # tangent_out is different, goes through the switch implementation.
  return switch_fun(arr, choice), switch_fun(arr_dot, choice)

good_custom_switchfun = jax.custom_jvp(switch_fun)
good_custom_switchfun.defjvp(good_switchfun_jvp)
print(jax.grad(good_custom_switchfun)(test_arr, test_val)) # 1. everywhere, good

bad_custom_switchfun = jax.custom_jvp(switch_fun)
bad_custom_switchfun.defjvp(bad_switchfun_jvp)
print(jax.grad(bad_custom_switchfun)(test_arr, test_val)) # ASSERTION ERROR

```


Discussion:
This variant of the repro works:

```python
import jax
import jax.numpy as jnp
jax.config.update('jax_enable_checks', True)


def mask_fun(arr, choice):
  out = (1 - choice) * arr.sum() +  choice * (1 - arr.sum())
  return out

def switch_fun(arr, choice):
  choice = jnp.floor(choice).astype(jnp.int32)
  out = jax.lax.switch(choice, [lambda x: x.sum(), lambda x: 1 - x.sum()], arr)
  return out

def switch_fun_linear(arr, choice):
  choice = jnp.floor(choice).astype(jnp.int32)
  from jax._src.lax.control_flow.conditionals import _cond
  out = _cond(choice, lambda x: 1 - x.sum(), lambda x: x.sum(), arr,
              linear=(True,))
  return out

test_arr = jnp.arange(3.)
test_val = 0.

print(mask_fun(test_arr, test_val))     # 3.0, all good
print(switch_fun(test_arr, test_val))   # 3.0, all good

print(jax.grad(mask_fun)(test_arr, test_val))   # -> 1. everywhere -> good
print(jax.grad(switch_fun)(test_arr, test_val)) # -> 1. everywhere -> good

def good_switchfun_jvp(primals, tangents):
  arr, choice = primals
  arr_dot, choice_dot = tangents
  return switch_fun(arr, choice), mask_fun(arr_dot, choice)

def bad_switchfun_jvp(primals, tangents):
  arr, choice = primals
  arr_dot, choice_dot = tangents
  # tangent_out is different, goes through the switch implementation.
  return switch_fun(arr, choice), switch_fun_linear(arr_dot, choice)

good_custom_switchfun = jax.custom_jvp(switch_fun)
good_custom_switchfun.defjvp(good_switchfun_jvp)
print(jax.grad(good_custom_switchfun)(test_arr, test_val)) # 1. everywhere, good

bad_custom_switchfun = jax.custom_jvp(switch_fun)
bad_custom_switchfun.defjvp(bad_switchfun_jvp)
print(jax.grad(bad_custom_switchfun)(test_arr, test_val)) # works now!
```

for control flow primitives like switch or cond, jax internally tracks in which inputs functions are guaranteed to be linear when the function is produced from jax's jvp autodiff pass. that linearity information is then useful when transposing.

originally (with scan, the first differentiable control flow operation) i think we only meant for that linearity information to be used in error-checking. but i think for cond we may have accidentally made it more load-bearing.

when using custom_jvp, rather than jax's internal jvp of switch, this linearity information isn't added by the jvp pass, and so downstream transposition fails

the above version of the repro works because i've made a switch_fun_linear variant of the switch_fun function to apply to the arr_dot values, which are linear inputs, and in switch_fun_linear i'm setting the linearity info explicitly. (i don't think the option is plumbed through to switch, so i used an internal cond.) this isn't the fix! manually annotating linearity information shouldn't be necessary. but i just wanted to test the hypothesis that somehow switch not getting linearity information was the problem

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
@@ -687,7 +687,9 @@ def transposed(*args):
   return _make_closed_jaxpr(transposed, res_avals + jaxpr.out_avals)
 
 def _cond_transpose(reduce_axes, cts, *args, branches, linear):
+  del linear  # could use for error checking, but see #14026
   index, *ops = args
+  linear = [type(x) is ad.UndefinedPrimal for x in ops]
   in_avals = map(raise_to_shaped, branches[0].in_avals)
   num_res = len(ops) - sum(linear)
   if any(isinstance(eff, state.RefEffect) for branch in branches for eff in


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
@@ -2564,5 +2564,45 @@ def body(c, _):
 
     jax.grad(f)(1.)  # doesn't crash
 
+  def test_custom_jvp_tangent_cond_transpose(self):
+    # https://github.com/google/jax/issues/14026
+    def mask_fun(arr, choice):
+      out = (1 - choice) * arr.sum() +  choice * (1 - arr.sum())
+      return out
+
+    def switch_fun(arr, choice):
+      choice = jnp.floor(choice).astype(jnp.int32)
+      out = jax.lax.switch(choice, [lambda x: x.sum(), lambda x: 1 - x.sum()], arr)
+      return out
+
+    test_arr = jnp.arange(3.)
+    test_val = 0.
+
+    expected1 = jax.grad(mask_fun)(test_arr, test_val)
+    expected2 = jax.grad(switch_fun)(test_arr, test_val)
+
+    def good_switchfun_jvp(primals, tangents):
+      arr, choice = primals
+      arr_dot, choice_dot = tangents
+      return switch_fun(arr, choice), mask_fun(arr_dot, choice)
+
+    def bad_switchfun_jvp(primals, tangents):
+      arr, choice = primals
+      arr_dot, choice_dot = tangents
+      return switch_fun(arr, choice), switch_fun(arr_dot, choice)
+
+    good_custom_switchfun = jax.custom_jvp(switch_fun)
+    good_custom_switchfun.defjvp(good_switchfun_jvp)
+    expected3 = jax.grad(good_custom_switchfun)(test_arr, test_val)
+
+    bad_custom_switchfun = jax.custom_jvp(switch_fun)
+    bad_custom_switchfun.defjvp(bad_switchfun_jvp)
+    actual = jax.grad(bad_custom_switchfun)(test_arr, test_val)
+
+    self.assertAllClose(expected1, expected2)
+    self.assertAllClose(expected2, expected3)
+    self.assertAllClose(expected3, actual)
+
+
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())