Repository: google/jax

Issue:
jax.lax.cond evaluates None as True
When using `jax.lax.cond`, passing `None` as the first argument causes the `true_fun` argument to get called. This is contrary to the behavior one would expect for Python if statements.


For example, running the following
```python
import jax

true_fun = lambda x: True
false_fun = lambda x: False

for pred in [True, False, None]:
	print(pred, jax.lax.cond(pred, true_fun, false_fun, None))
```

gives the following output:
```
True True
False False
None True
```

I am using the newest PyPI versions of `jax` and `jaxlib` (0.3.14).


Discussion:
For what it's worth, the reason `None` evaluates to `True` is because of this:
```python
 >>> jnp.array(None)
DeviceArray(nan, dtype=float32)
```
And `NaN`, when converted to `bool`, is not False because it's not equal to zero:
```python
>>> jnp.array(None).astype(bool)                                                                                                                                                                                                                                                                      
DeviceArray(True, dtype=bool)
```
I suppose a better behavior here would be to raise an error when `None` is passed in a place where a JAX array is expected.
That would work. The only thing to keep an eye on is the documentation of `jax.lax.cond`, which currently says:

```
cond() has equivalent semantics to this Python implementation:

  def cond(pred, true_fun, false_fun, *operands):
    if pred:
      return true_fun(*operands)
    else:
      return false_fun(*operands)
```

Since `bool(None)` is `False`, some caveat would be helpful to warn users about unexpected errors that wouldn't arise in the Python equivalent.

Regardless, throwing an error is still better than the present situation, since the branch that is selected is currently the opposite of what the documentation implies.

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
@@ -148,7 +148,8 @@ def _cond(pred, true_fun: Callable, false_fun: Callable, *operands,
           operand=_no_operand_sentinel, linear=None):
   """Conditionally apply ``true_fun`` or ``false_fun``.
 
-  ``cond()`` has equivalent semantics to this Python implementation::
+  Provided arguments are correctly typed, ``cond()`` has equivalent
+  semantics to this Python implementation::
 
     def cond(pred, true_fun, false_fun, *operands):
       if pred:
@@ -181,6 +182,8 @@ def cond(pred, true_fun, false_fun, *operands):
     operands = (operand,)
   del operand
 
+  if pred is None:
+    raise TypeError("cond predicate is None")
   if isinstance(pred, Sequence) or np.ndim(pred) != 0:
     raise TypeError(
         f"Pred must be a scalar, got {pred} of " +


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
@@ -628,6 +628,16 @@ def false_fun(x):
     self.assertEqual(fun(4), cfun(4))
     self.assertEqual(fun(4), (8, 16))
 
+  def testCondPredIsNone(self):
+    # see https://github.com/google/jax/issues/11574
+    def f(pred, x):
+      return lax.cond(pred, lambda x: x + 1, lambda x: x + 2, x)
+
+    self.assertRaisesRegex(TypeError, "cond predicate is None",
+                           lambda: f(None, 1.))
+    self.assertRaisesRegex(TypeError, "cond predicate is None",
+                           lambda: jax.jit(f)(None, 1.))
+
   def testCondTwoOperands(self):
     # see https://github.com/google/jax/issues/8469
     add, mul = lax.add, lax.mul