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


Discussion:
hm, I think this means we need a `switch_error_check` which will do this functionalizing/merging for us. Thanks for reporting! (and for bearing with us while we make checkify feature complete!)
aha, actually the issue is that the `lax.switch` doesn't have any inputs: if you add an operand, the check is succesfully eliminated (so you don't get the first functionalization error).

```
def make_branch(i):                                                                                                                                      
  def branch(x):                                                                                                                                         
    checkify.check(jnp.sum(x) > 0, 'Failed!')                                                                                                            
    return x                                                                                                                                             
  return branch                                                                                                                                          
                                                                                                                                                         
def f(x):                                                                                                                                                
  return lax.switch(0, [make_branch(i) for i in range(3)], x)                                                                                            
                                                                                                                                                         
checked_f = checkify.checkify(f)                                                                                                                         
                                                                                                                                                         
err, _ = checked_f(-1.)                                                                                                                                  
```

This still needs to be fixed, but might unblock you for the moment?

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
@@ -139,6 +139,10 @@ def switch(index, branches, *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, false_fun, *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(


Test patch:
diff --git a/tests/checkify_test.py b/tests/checkify_test.py
--- a/tests/checkify_test.py
+++ b/tests/checkify_test.py
@@ -766,6 +766,33 @@ def f(init_i, init_val):
     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 @@ def g(x):
     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 @@ def while_body(s):
     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 @@ def f(x):
     with self.assertRaisesRegex(xla_extension.XlaRuntimeError,
                                 "x needs to be positive"):
       f(-1.)
-
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())