Repository: google/jax
Issue: Problem mixing checkify.check and jax.lax.switch

### Description

I'm trying to do a switch over functions that have a checkify check inside. It goes somewhat like this:
```python
def make_branch(i):
  def branch():
    result = jnp.full((1,), i)
    checkify.check(jnp.sum(result) > 0, 'Failed!')
    return result
  return branch

index = jnp.full((), 0)
result = jax.lax.switch(index, [make_branch(i) for i in range(3)])
```
This, unfortunately, fails with `Cannot abstractly evaluate a checkify.check which was not functionalized`, presumably because `switch` compiles its branches.

I've tried transforming branches with `checkify`:
```python
def make_branch(i):
  @checkify.checkify
  def branch():
    result = jnp.full((1,), i)
    checkify.check(jnp.sum(result) > 0, 'Failed!')
    return result
  return branch

index = jnp.full((), 0)
err, result = jax.lax.switch(index, [make_branch(i) for i in range(3)])
checkify.check_error(err)
```
This approach, unfortunately, fails with another error:
```
branch 0 and 1 outputs must have same type structure, got PyTreeDef((CustomNode(Error[((4, 'Failed! (check failed at <ipython-input-14-15495f6829f6>:6 (branch))'),)], [*, *, *]), *)) and PyTreeDef((CustomNode(Error[((5, 'Failed! (check failed at <ipython-input-14-15495f6829f6>:6 (branch))'),)], [*, *, *]), *)).
```
For some reason type signatures for branches are different where error is envolved. That seems like an unexpected behavior that prohibits using checks with conditional constructs. Is there a way around it?

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

jax v0.3.17, jaxlib v0.3.15

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

GPU

### Additional System Info

Google Colab

Changed files:
tests/checkify_test.py, jax/_src/lax/control_flow/conditionals.py

diff --git a/jax/_src/lax/control_flow/conditionals.py b/jax/_src/lax/control_flow/conditionals.py
index ed82f9807..928875e70 100644
--- a/jax/_src/lax/control_flow/conditionals.py
+++ b/jax/_src/lax/control_flow/conditionals.py
@@ -139,6 +139,10 @@ def switch(index, branches: Sequence[Callable], *operands,
   if disallowed_effects:
     raise NotImplementedError(
         f'Effects not supported in `switch`: {disallowed_effects}')
+  if joined_effects:
+    # Raise index in case of effects to allow data-dependence-based discharging
+    # of those effects (even if they don't have an explicit data dependence).
+    index = core.raise_as_much_as_possible(index)
 
   linear = (False,) * (len(consts) + len(ops))
   out = cond_p.bind(
@@ -235,6 +239,10 @@ def _cond(pred, true_fun: Callable, false_fun: Callable, *operands,
         f'Effects not supported in `cond`: {disallowed_effects}')
 
   index = lax.convert_element_type(pred, np.int32)
+  if joined_effects:
+    # Raise index in case of effects to allow data-dependence-based discharging
+    # of those effects (even if they don't have an explicit data dependence).
+    index = core.raise_as_much_as_possible(index)
 
   linear = [False] * len(consts) + linear_ops
   out = cond_p.bind(
diff --git a/tests/checkify_test.py b/tests/checkify_test.py
index 0866957e2..2afe3a254 100644
--- a/tests/checkify_test.py
+++ b/tests/checkify_test.py
@@ -766,6 +766,33 @@ class AssertPrimitiveTests(jtu.JaxTestCase):
     self.assertIsNotNone(err.get())
     self.assertStartsWith(err.get(), "x must be negative")
 
+  def test_assert_discharging_cond(self):
+    def true_branch(x):
+      checkify.check(jnp.all(x != 0.), "x cannot be 0")
+      return 1/x
+
+    def false_branch(x):
+      checkify.check(jnp.all(x >= 0), "x must be positive")
+      return x*2
+
+    @jax.jit
+    def f(pred, x):
+      return lax.cond(pred, true_branch, false_branch, x)
+
+    checked_f = checkify.checkify(f)
+
+    err, _ = checked_f(True, 0.)
+    self.assertIsNotNone(err.get())
+    self.assertStartsWith(err.get(), "x cannot be 0")
+    err, _ = checked_f(False, 0.)
+    self.assertIsNone(err.get())
+
+    err, _ = checked_f(False, -1.)
+    self.assertIsNotNone(err.get())
+    self.assertStartsWith(err.get(), "x must be positive")
+    err, _ = checked_f(True, -1.)
+    self.assertIsNone(err.get())
+
   def test_assert_batching_rule(self):
     @jax.vmap
     def f(x):
@@ -863,12 +890,12 @@ class AssertPrimitiveTests(jtu.JaxTestCase):
     checkify.checkify(g)(0.)  # does not crash
 
   def test_grad(self):
-    @checkify.checkify
     @jax.grad
     def f(x):
       checkify.check(jnp.all(x > 0), "should be positive!")
       return x
 
+    f = checkify.checkify(f)
     err, _ = f(1.)
     self.assertIsNone(err.get())
 
@@ -906,6 +933,30 @@ class AssertPrimitiveTests(jtu.JaxTestCase):
     self.assertIsNotNone(err.get())
     self.assertStartsWith(err.get(), "value needs to be less than 6")
 
+  def test_assert_cond_no_data_dependence(self):
+    def f():
+      return jax.lax.cond(True,
+                          lambda: checkify.check(False, "hi!"),
+                          lambda: checkify.check(False, "bye!"))
+
+    f = checkify.checkify(f)
+    err, _ = f()
+    self.assertIsNotNone(err.get())
+    self.assertStartsWith(err.get(), "hi!")
+
+  def test_assert_switch_no_data_dependence(self):
+    def branch():
+      checkify.check(False, "hi!")
+
+    def f():
+      return lax.switch(0, [branch]*3)
+
+    checked_f = checkify.checkify(f)
+
+    err, _ = checked_f()
+    self.assertIsNotNone(err.get())
+    self.assertStartsWith(err.get(), "hi!")
+
 class LowerableChecksTest(jtu.JaxTestCase):
   def setUp(self):
     super().setUp()
@@ -926,6 +977,5 @@ class LowerableChecksTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(xla_extension.XlaRuntimeError,
                                 "x needs to be positive"):
       f(-1.)
-
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())