Repository: google/jax

Issue:
lax.cond should not fail with RuntimeError when passed a integer
It would be reasonable to convert integer predicates into booleans by checking if they are equal to zero, like how Python handles integers.

Example behavior:
```
import jax
from jax import lax

def f(n):
  return lax.cond(n % 2, n, lambda n: n, n, lambda n: n)

f(12)
```
Results in:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-17-5bf9fbe6fc46> in <module>()
      5   return lax.cond(n % 2, n, lambda n: n, n, lambda n: n)
      6 
----> 7 f(12)

5 frames
<ipython-input-17-5bf9fbe6fc46> in f(n)
      3 
      4 def f(n):
----> 5   return lax.cond(n % 2, n, lambda n: n, n, lambda n: n)
      6 
      7 f(12)

google3/third_party/py/jax/lax/lax_control_flow.py in cond(pred, true_operand, true_fun, false_operand, false_fun)
    295       *itertools.chain([pred], true_consts, true_ops, false_consts, false_ops),
    296       true_jaxpr=true_jaxpr, false_jaxpr=false_jaxpr,
--> 297       true_nconsts=len(true_consts), false_nconsts=len(false_consts))
    298   return tree_unflatten(out_tree, out)
    299 

google3/third_party/py/jax/core.py in bind(self, *args, **kwargs)
    128     top_trace = find_top_trace(args)
    129     if top_trace is None:
--> 130       return self.impl(*args, **kwargs)
    131 
    132     tracers = map(top_trace.full_raise, args)

google3/third_party/py/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    121   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
    122   abstract_args = map(abstractify, args)
--> 123   compiled_fun = xla_primitive_callable(prim, *abstract_args, **params)
    124   return compiled_fun(*args)
    125 

google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *abstract_args, **params)
    134     handle_result = aval_to_result_handler(aval_out)
    135   xla_shapes = tuple(map(aval_to_xla_shape, abstract_args))
--> 136   built_c = primitive_computation(prim, *xla_shapes, **params)
    137   compiled = built_c.Compile(xla_shapes, xb.get_compile_options(),
    138                              backend=xb.get_backend(backend))

google3/third_party/py/jax/interpreters/xla.py in primitive_computation(prim, *xla_shapes, **params)
    166            "This is a bug in JAX's shape-checking rules; please report it!\n"
    167            "https://github.com/google/jax/issues\n")
--> 168     raise RuntimeError(msg)
    169 
    170 def _execute_compiled_primitive(prim, compiled, backend, result_handler, *args):

RuntimeError: Invalid argument: Argument to predicated-Conditional is not a scalar of PRED type (s32[]).:     @     0x55e6131c5b9d  xla::Conditional()
    @     0x7f0e34c7e9a1  (unknown)
    @     0x7f0e34c91e6b  (unknown)
    @     0x55e621de056d  _PyCFunction_FastCallDict
    @     0x55e6216c1a82  _PyEval_EvalFrameDefault
    @     0x55e6216ba7d0  _PyEval_EvalFrameDefault
    @     0x55e6210eb7ce  _PyEval_EvalCodeWithName.llvm.293474707488941182
    @     0x55e62546cd7c  function_call
    @     0x55e6216c1564  _PyEval_EvalFrameDefault
    @     0x55e6210eb7ce  _PyEval_EvalCodeWithName.llvm.293474707488941182
    @     0x55e62546cd7c  function_call
    @     0x55e6210dfdaf  PyObject_Call
    @     0x7f0e349bdbda  (unknown)
    @     0x55e6216c1564  _PyEval_EvalFrameDefault
    @     0x55e6210eb7ce  _PyEval_EvalCodeWithName.llvm.293474707488941182
    @     0x55e62546cd7c  function_call
    @     0x55e6210dfdaf  PyObject_Call
    @     0x7f0e349bdbda  (unknown)
    @     0x55e6216c1564  _PyEval_EvalFrameDefault
    @     0x55e6210eb7ce  _PyEval_EvalCodeWithName.llvm.293474707488941182
    @     0x55e62546cd7c  function_call
    @     0x55e6210dfdaf  PyObject_Call
    @     0x55e6240cd6ed  partial_call
    @     0x55e6216c1564  _PyEval_EvalFrameDefault
    @     0x55e6210eb7ce  _PyEval_EvalCodeWithName.llvm.293474707488941182
    @     0x55e620927e47  _PyFunction_FastCallDict
    @     0x55e6254651b8  method_call
    @     0x55e6216c1564  _PyEval_EvalFrameDefault
    @     0x55e6216ba7d0  _PyEval_EvalFrameDefault
    @     0x55e6216ba7d0  _PyEval_EvalFrameDefault
    @     0x55e6210eb7ce  _PyEval_EvalCodeWithName.llvm.293474707488941182
    @     0x55e61c77b679  builtin_exec

This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues
```

The `RuntimeError` here is particularly problematic because it mean that if you `jit` such a function, you don't get a meaningful traceback:
```
>>> jax.jit(f)(12)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-18-05daf300cab6> in <module>()
----> 1 jax.jit(f)(12)

9 frames
google3/third_party/py/jax/api.py in f_jitted(*args, **kwargs)
    147     _check_args(args_flat)
    148     flat_fun, out_tree = flatten_fun(f, in_tree)
--> 149     out = xla.xla_call(flat_fun, *args_flat, device_assignment=device_assignment, backend=backend)
    150     return tree_unflatten(out_tree(), out)
    151 

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **params)
    567   if top_trace is None:
    568     with new_sublevel():
--> 569       outs = primitive.impl(f, *args, **params)
    570   else:
    571     tracers = map(top_trace.full_raise, args)

google3/third_party/py/jax/interpreters/xla.py in _xla_call_impl(fun, *args, **params)
    365   backend = params.get('backend', None)
    366   compiled_fun = _xla_callable(fun, device_assignment, backend,
--> 367                                *map(abstractify, args))
    368   try:
    369     return compiled_fun(*args)

google3/third_party/py/jax/linear_util.py in cached_fun(f, *args)
    215 
    216   def cached_fun(f, *args):
--> 217     ans, f_prev = cached_fun_body(f, args)
    218     if id(f_prev) != id(f):
    219       f.populate_stores(f_prev)

google3/third_party/py/jax/linear_util.py in cached_fun_body(f, args)
    212   @fastcache.clru_cache(maxsize=max_size)
    213   def cached_fun_body(f, args):
--> 214     return call(f, *args), f
    215 
    216   def cached_fun(f, *args):

google3/third_party/py/jax/interpreters/xla.py in _xla_callable(fun, device_assignment, backend, *abstract_args)
    383     axis_env = AxisEnv(jaxpr_replicas(jaxpr), [], [])
    384     compiled = compile_jaxpr(jaxpr, device_assignment, backend, axis_env, consts,
--> 385                              *abstract_args)
    386     del master, consts, jaxpr, env
    387   result_handlers = tuple(map(_pval_to_result_handler, pvals))

google3/third_party/py/jax/interpreters/xla.py in compile_jaxpr(jaxpr, device_assignment, backend, axis_env, const_vals, *abstract_args)
    200     raise ValueError(msg.format(axis_env.nreps, xb.device_count(backend)))
    201   arg_shapes = tuple(map(aval_to_xla_shape, abstract_args))
--> 202   built_c = jaxpr_computation(jaxpr, backend, axis_env, const_vals, (), *arg_shapes)
    203   compile_opts = xb.get_compile_options(num_replicas=axis_env.nreps,
    204                                         device_assignment=device_assignment)

google3/third_party/py/jax/interpreters/xla.py in jaxpr_computation(jaxpr, backend, axis_env, const_vals, freevar_shapes, *arg_shapes)
    239   freevars = _map(c.ParameterWithShape, freevar_shapes)
    240   args = _map(c.ParameterWithShape, arg_shapes)
--> 241   out_nodes = jaxpr_subcomp(c, jaxpr, backend, axis_env, consts, freevars, *args)
    242   return c.Build(c.Tuple(*out_nodes))
    243 

google3/third_party/py/jax/interpreters/xla.py in jaxpr_subcomp(c, jaxpr, backend, axis_env, consts, freevars, *args)
    292       raise NotImplementedError(msg.format(eqn.primitive.name))
    293 
--> 294     c.GetShape(ans)  # force xla to do shape error checking
    295     out_nodes = xla_destructure(c, ans) if eqn.primitive.multiple_results else [ans]
    296     _map(write, eqn.outvars, out_nodes)

google3/third_party/tensorflow/compiler/xla/python/xla_client.py in GetShape(self, operand)
    718 
    719   def GetShape(self, operand):
--> 720     return self._builder.GetShape(operand)
    721 
    722   def SetOpMetadata(self, op_metadata):

RuntimeError: Invalid argument: Argument to predicated-Conditional is not a scalar of PRED type (s32[]).
```


Discussion:
Is this just adding a 
```python
if isinstance(pred, int):
    pred = pred != 0
```
in the [cond](https://github.com/google/jax/blob/master/jax/lax/lax_control_flow.py#L279) function? If yes I'm happy to submit a PR
I think `pred` probably should be converted into an array first and this conversion should be done on all integer dtypes. Otherwise this won't work inside `jit`.
I get the handling all int dtypes for `jit` part; converting to array is because we are calling `itertools.chain`?

I'll try to have a go and add some tests as well.
Converting to an array is just to ensure `pred` is in a consistent form with `pred.dtype` defined. Otherwise this would fail when you actually pass a Python integer is as `pred`, as in my example above, because Python integers don't have a dtype.

Resolving patch:
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -286,7 +286,21 @@ def cond(pred, true_operand, true_fun, false_operand, false_fun):
         return true_fun(true_operand)
       else:
         return false_fun(false_operand)
+
+  Pred has to be a scalar type, collection types (list, tuple) are not supported
+
   """
+
+  if len(onp.shape(pred)) != 0:
+    raise TypeError("Pred must be a scalar, got {} of shape {}".format(pred, onp.shape(pred)))
+
+  pred_dtype = onp.result_type(pred)
+  if pred_dtype.kind != 'b':
+    if pred_dtype.kind in 'iuf':
+      pred = pred != 0
+    else:
+      msg = ("Pred type must be either boolean or number, got {}")
+      raise TypeError(msg.format(pred_dtype))
   true_ops, true_tree = tree_flatten((true_operand,))
   true_avals = tuple(_map(_abstractify, true_ops))
   true_jaxpr, true_consts, out_tree = _initial_style_jaxpr(true_fun, true_tree, true_avals)


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
@@ -416,6 +416,25 @@ def false_fun(x):
     self.assertEqual(fun(4), cfun(4))
     self.assertEqual(fun(4), (8, 16))
 
+  def testIssue1379(self):
+
+    def fun(pred):
+      return lax.cond(pred, pred, lambda x: (True, x), pred, lambda x: (False, x))
+    
+    @api.jit
+    def cfun(pred):
+      return fun(pred)
+    
+    self.assertEqual(fun(0), cfun(0), (False,0))
+    self.assertEqual(fun(0.), cfun(0.), (False,0.))
+    self.assertEqual(fun(1), cfun(1), (True,1))
+    self.assertEqual(fun(1.), cfun(1.), (True,1.))
+
+    # test that proper errors are raised for wrong types
+    for pred in ["abc", [], [1,2]]:
+      for f in [fun, cfun]:
+        self.assertRaises(TypeError, f, pred)
+    
   def testNestedCond(self):
     def fun(x):
       if x < 2: