Repository: google/jax

Issue:
einsum crash on bool dtype
In colab,

```python
import jax, jax.numpy as jp
x = jp.array(np.ones((3, 3), dtype=bool))
print(jp.einsum("ij->i", x))
```

results in

> RuntimeError: Invalid argument: Expected element type in shape to be arithmetic type for operation add; got PRED.: 
> This is a bug in JAX's shape-checking rules; please report it!
> https://github.com/google/jax/issues

Note that in numpy, `einsum` on `bool` dtypes does a logical OR reduction ("einany") rather than summation.


Discussion:

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -535,18 +535,6 @@ def dot(lhs, rhs, precision=None):
   Returns:
     An array containing the product.
   """
-  # TODO(b/134526360): XLA doesn't support integer dots, so we emit a sum of
-  # products instead.
-  if dtypes.issubdtype(lhs.dtype, onp.integer):
-    lhs_shape = onp.shape(lhs)
-    lhs_ndim = len(lhs_shape)
-    rhs_ndim = onp.ndim(rhs)
-    if rhs_ndim > 1:
-      lhs = broadcast_in_dim(lhs, lhs_shape + (1,), tuple(range(len(lhs_shape))))
-    if lhs_ndim > 1:
-      rhs = broadcast(rhs, (1,))
-    return reduce(mul(lhs, rhs), _zero(lhs), add, (len(lhs_shape) - 1,))
-
   if 1 <= lhs.ndim <= 2 and 1 <= rhs.ndim <= 2 and lhs.shape[-1] == rhs.shape[0]:
     return dot_general(lhs, rhs, (((lhs.ndim - 1,), (0,)), ((), ())),
                        precision=precision)
@@ -576,9 +564,9 @@ def dot_general(lhs, rhs, dimension_numbers, precision=None):
   contract_dims, batch_dims = dimension_numbers
   contract_dims = tuple(map(tuple, contract_dims))
   batch_dims = tuple(map(tuple, batch_dims))
-  if dtypes.issubdtype(lhs.dtype, onp.integer):
-    # TODO(b/134526360): XLA doesn't support integer dots, so we emit a sum of
-    # products instead.
+  if not dtypes.issubdtype(lhs.dtype, onp.inexact):
+    # TODO(b/134526360): XLA doesn't support bool or integer dots, so we emit a
+    # sum of products instead.
     lhs_contract_dims, rhs_contract_dims = contract_dims
     lhs_batch_dims, rhs_batch_dims = batch_dims
     lhs_noncontract_dims = tuple(sorted(
@@ -589,16 +577,19 @@ def dot_general(lhs, rhs, dimension_numbers, precision=None):
                     lhs_batch_dims + lhs_noncontract_dims + lhs_contract_dims)
     rhs = transpose(rhs,
                     rhs_batch_dims + rhs_noncontract_dims + rhs_contract_dims)
-    new_lhs_shape = onp.insert(
-      onp.shape(lhs), len(lhs_batch_dims) + len(lhs_noncontract_dims),
-      (1,) * len(rhs_noncontract_dims))
-    new_rhs_shape = onp.insert(onp.shape(rhs), len(lhs_batch_dims),
+    new_lhs_shape = onp.insert(onp.array(onp.shape(lhs), dtype=onp.int64),
+                               len(lhs_batch_dims) + len(lhs_noncontract_dims),
+                               (1,) * len(rhs_noncontract_dims))
+    new_rhs_shape = onp.insert(onp.array(onp.shape(rhs), dtype=onp.int64),
+                               len(lhs_batch_dims),
                                (1,) * len(lhs_noncontract_dims))
     lhs = reshape(lhs, new_lhs_shape)
     rhs = reshape(rhs, new_rhs_shape)
     out_ndim = (len(lhs_batch_dims) + len(lhs_noncontract_dims) +
                 len(rhs_noncontract_dims))
-    return reduce(mul(lhs, rhs), _zero(lhs), add,
+    op_product = bitwise_and if lhs.dtype == onp.bool_ else mul
+    op_sum = bitwise_or if lhs.dtype == onp.bool_ else add
+    return reduce(op_product(lhs, rhs), _zero(lhs), op_sum,
                   tuple(range(out_ndim, out_ndim + len(lhs_contract_dims))))
 
   return dot_general_p.bind(lhs, rhs,
@@ -1526,7 +1517,7 @@ def unop_dtype_rule(result_dtype, accepted_dtypes, name, aval, **kwargs):
   if not any(dtypes.issubdtype(aval.dtype, t) for t in accepted_dtypes):
     msg = '{} does not accept dtype {}. Accepted dtypes are subtypes of {}.'
     typename = str(onp.dtype(aval.dtype).name)
-    accepted_typenames = (str(onp.dtype(t).name) for t in accepted_dtypes)
+    accepted_typenames = (t.__name__ for t in accepted_dtypes)
     raise TypeError(msg.format(name, typename, ', '.join(accepted_typenames)))
   return result_dtype(aval.dtype)
 
@@ -1548,7 +1539,7 @@ def binop_dtype_rule(result_dtype, accepted_dtypes, name, *avals, **kwargs):
       msg = ('{} does not accept dtype {} at position {}. '
              'Accepted dtypes at position {} are subtypes of {}.')
       typename = str(onp.dtype(aval_dtype).name)
-      typenames = ', '.join(str(onp.dtype(t).name) for t in types)
+      typenames = ', '.join(t.__name__ for t in types)
       raise TypeError(msg.format(name, typename, i, i, typenames))
   _check_same_dtypes(name, False, *aval_dtypes)
   return result_dtype(*avals)
@@ -1611,7 +1602,7 @@ def _brcast_to(x, shape):
     return broadcast(x, shape)
 
 
-_float = {onp.floating, dtypes.bfloat16}
+_float = {onp.floating}
 _complex = {onp.complexfloating}
 _complex_elem_types = {onp.float32, onp.float64}
 _int = {onp.integer}
@@ -3384,6 +3375,12 @@ def _reducer_masking_rule(prim, identity, padded_vals, logical_shapes,
 batching.primitive_batchers[reduce_p] = _reduce_batch_rule
 
 
+def _reduce_number_dtype_rule(name, operand, *args, **kw):
+  if not dtypes.issubdtype(operand.dtype, onp.number):
+    raise TypeError("{} does not accept dtype {}. Accepted dtypes are subtypes "
+                    "of number.".format(name, onp.dtype(operand.dtype).name))
+  return dtypes.canonicalize_dtype(operand.dtype)
+
 def _reduce_sum_shape_rule(operand, axes, input_shape):
   assert operand.shape == input_shape, ('{} != {}'
                                         .format(operand.shape, input_shape))
@@ -3402,8 +3399,9 @@ def _reduce_sum_transpose_rule(cotangent, input_shape, axes):
   assert result.shape == input_shape
   return [result]
 
-reduce_sum_p = standard_primitive(_reduce_sum_shape_rule, _input_dtype,
-                                  'reduce_sum', _reduce_sum_translation_rule)
+reduce_sum_p = standard_primitive(
+  _reduce_sum_shape_rule, partial(_reduce_number_dtype_rule, 'reduce_sum'),
+  'reduce_sum', _reduce_sum_translation_rule)
 ad.deflinear(reduce_sum_p, _reduce_sum_transpose_rule)
 batching.defreducer(reduce_sum_p)
 _masking_defreducer(reduce_sum_p,
@@ -3450,8 +3448,9 @@ def _reduce_prod_jvp_rule(tangent, operand, axes):
   # Multiply partial products with the tangents and sum.
   return _reduce_sum(mul(tangent, mul(left_products, right_products)), (0,))
 
-reduce_prod_p = standard_primitive(_reduce_prod_shape_rule, _input_dtype,
-                                   'reduce_prod', _reduce_prod_translation_rule)
+reduce_prod_p = standard_primitive(
+  _reduce_prod_shape_rule, partial(_reduce_number_dtype_rule, 'reduce_prod'),
+  'reduce_prod', _reduce_prod_translation_rule)
 ad.defjvp(reduce_prod_p, _reduce_prod_jvp_rule)
 batching.defreducer(reduce_prod_p)
 
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -388,6 +388,12 @@ def _one_to_one_binop(numpy_fn, lax_fn, promote_to_inexact=False):
     fn = lambda x1, x2: lax_fn(*_promote_args(numpy_fn.__name__, x1, x2))
   return _wraps(numpy_fn)(fn)
 
+def _maybe_bool_binop(numpy_fn, lax_fn, bool_lax_fn):
+  def fn(x1, x2):
+    x1, x2 = _promote_args(numpy_fn.__name__, x1, x2)
+    return lax_fn(x1, x2) if x1.dtype != bool_ else bool_lax_fn(x1, x2)
+  return _wraps(numpy_fn)(fn)
+
 absolute = abs = _one_to_one_unop(onp.absolute, lax.abs)
 fabs = _one_to_one_unop(onp.fabs, lax.abs, True)
 bitwise_not = _one_to_one_unop(onp.bitwise_not, lax.bitwise_not)
@@ -413,14 +419,14 @@ def _one_to_one_binop(numpy_fn, lax_fn, promote_to_inexact=False):
 sqrt = _one_to_one_unop(onp.sqrt, lax.sqrt, True)
 
 
-add = _one_to_one_binop(onp.add, lax.add)
+add = _maybe_bool_binop(onp.add, lax.add, lax.bitwise_or)
 bitwise_and = _one_to_one_binop(onp.bitwise_and, lax.bitwise_and)
 bitwise_or = _one_to_one_binop(onp.bitwise_or, lax.bitwise_or)
 bitwise_xor = _one_to_one_binop(onp.bitwise_xor, lax.bitwise_xor)
 right_shift = _one_to_one_binop(onp.right_shift, lax.shift_right_arithmetic)
 left_shift = _one_to_one_binop(onp.left_shift, lax.shift_left)
 equal = _one_to_one_binop(onp.equal, lax.eq)
-multiply = _one_to_one_binop(onp.multiply, lax.mul)
+multiply = _maybe_bool_binop(onp.multiply, lax.mul, lax.bitwise_and)
 not_equal = _one_to_one_binop(onp.not_equal, lax.ne)
 subtract = _one_to_one_binop(onp.subtract, lax.sub)
 arctan2 = _one_to_one_binop(onp.arctan2, lax.atan2, True)
@@ -2244,7 +2250,9 @@ def einsum_path(subscripts, *operands, **kwargs):
 @partial(jit, static_argnums=(1, 2))
 def _einsum(operands, contractions, precision):
   operands = list(_promote_dtypes(*operands))
-  sum = lambda x, axes: lax.reduce(x, onp.array(0, x.dtype), lax.add, axes)
+  def sum(x, axes):
+    return lax.reduce(x, onp.array(0, x.dtype),
+                      lax.add if x.dtype != bool_ else lax.bitwise_or, axes)
 
   def sum_uniques(operand, names, uniques):
     if uniques:


Test patch:
diff --git a/tests/lax_numpy_einsum_test.py b/tests/lax_numpy_einsum_test.py
--- a/tests/lax_numpy_einsum_test.py
+++ b/tests/lax_numpy_einsum_test.py
@@ -225,7 +225,8 @@ def test_tf_unsupported_3(self):
 
   # these tests are based on https://github.com/dask/dask/pull/3412/files
   @parameterized.named_parameters(
-      {"testcase_name": "_{}".format(einstr), "einstr": einstr}
+      {"testcase_name": "_{}_dtype={}".format(einstr, dtype.__name__),
+      "einstr": einstr, "dtype": dtype}
       for einstr in [
           'abc,bad->abcd',
           'abcdef,bcdfg->abcdeg',
@@ -256,9 +257,10 @@ def test_tf_unsupported_3(self):
           'aab,bcc->ac',
           'fdf,cdd,ccd,afe->ae',
           'fff,fae,bef,def->abd',
-      ])
-  def test_from_dask(self, einstr):
-    r = rng()
+      ]
+      for dtype in [np.float32, np.int32, np.complex64, np.bool_])
+  def test_from_dask(self, einstr, dtype):
+    r = jtu.rand_default()
     if '->' in einstr:
       input_str, result_names = einstr.split('->')
     else:
@@ -269,7 +271,7 @@ def test_from_dask(self, einstr):
     shapes = defaultdict(lambda: next(dims))
     input_shapes = [tuple(shapes[c] for c in names.replace('...', '01'))
                     for names in input_names]
-    operands = [r.randn(*shape) for shape in input_shapes]
+    operands = [r(shape, dtype) for shape in input_shapes]
 
     self._check(einstr, *operands)
 
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -96,7 +96,7 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
 
 JAX_ONE_TO_ONE_OP_RECORDS = [
     op_record("abs", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
-    op_record("add", 2, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
+    op_record("add", 2, all_dtypes, all_shapes, jtu.rand_default, ["rev"]),
     op_record("ceil", 1, float_dtypes, all_shapes, jtu.rand_default, []),
     op_record("conj", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
     op_record("equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
@@ -109,23 +109,23 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
                          onp.float64: 1e-12, onp.complex64: 2e-4,
                          onp.complex128: 1e-12}, check_dtypes=False),
     op_record("floor", 1, float_dtypes, all_shapes, jtu.rand_default, []),
-    op_record("greater", 2, number_dtypes, all_shapes, jtu.rand_some_equal, []),
-    op_record("greater_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal, []),
-    op_record("less", 2, number_dtypes, all_shapes, jtu.rand_some_equal, []),
-    op_record("less_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal, []),
+    op_record("greater", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
+    op_record("greater_equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
+    op_record("less", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
+    op_record("less_equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal, []),
     op_record("log", 1, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],
               inexact=True),
     op_record("logical_and", 2, all_dtypes, all_shapes, jtu.rand_bool, []),
     op_record("logical_not", 1, all_dtypes, all_shapes, jtu.rand_bool, []),
     op_record("logical_or", 2, all_dtypes, all_shapes, jtu.rand_bool, []),
     op_record("logical_xor", 2, all_dtypes, all_shapes, jtu.rand_bool, []),
-    op_record("maximum", 2, number_dtypes, all_shapes, jtu.rand_some_inf, []),
-    op_record("minimum", 2, number_dtypes, all_shapes, jtu.rand_some_inf, []),
-    op_record("multiply", 2, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
+    op_record("maximum", 2, all_dtypes, all_shapes, jtu.rand_some_inf, []),
+    op_record("minimum", 2, all_dtypes, all_shapes, jtu.rand_some_inf, []),
+    op_record("multiply", 2, all_dtypes, all_shapes, jtu.rand_default, ["rev"]),
     op_record("negative", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
     op_record("nextafter", 2, [f for f in float_dtypes if f != lnp.bfloat16],
               all_shapes, jtu.rand_default, ["rev"], inexact=True, tolerance=0),
-    op_record("not_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal, ["rev"]),
+    op_record("not_equal", 2, all_dtypes, all_shapes, jtu.rand_some_equal, ["rev"]),
     op_record("array_equal", 2, number_dtypes, all_shapes, jtu.rand_some_equal, ["rev"]),
     op_record("reciprocal", 1, inexact_dtypes, all_shapes, jtu.rand_default, []),
     op_record("subtract", 2, number_dtypes, all_shapes, jtu.rand_default, ["rev"]),
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -634,7 +634,7 @@ def fun_via_grad(lhs, rhs):
        "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
        "precision": precision, "rng_factory": rng_factory}
       for lhs_shape in [(3,), (4, 3)] for rhs_shape in [(3,), (3, 6)]
-      for dtype in default_dtypes
+      for dtype in all_dtypes
       for precision in [None, lax.Precision.DEFAULT, lax.Precision.HIGH,
                         lax.Precision.HIGHEST]
       for rng_factory in [jtu.rand_default]))
@@ -651,14 +651,16 @@ def testDot(self, lhs_shape, rhs_shape, dtype, precision, rng_factory):
        "lhs_shape": lhs_shape, "rhs_shape": rhs_shape, "dtype": dtype,
        "rng_factory": rng_factory}
       for lhs_shape in [(3,), (4, 3)] for rhs_shape in [(3,), (3, 6)]
-      for dtype in default_dtypes
+      for dtype in all_dtypes
       for rng_factory in [jtu.rand_default]))
   def testDotAgainstNumpy(self, lhs_shape, rhs_shape, dtype, rng_factory):
     rng = rng_factory()
     args_maker = lambda: [rng(lhs_shape, dtype), rng(rhs_shape, dtype)]
     tol = {
       onp.float16: 1e-2,
-      onp.float64: max(jtu.default_tolerance()[onp.dtype(onp.float64)], 1e-14)
+      onp.float64: max(jtu.default_tolerance()[onp.dtype(onp.float64)], 1e-14),
+      onp.complex128: max(jtu.default_tolerance()[onp.dtype(onp.complex128)],
+                          1e-14)
     }
     lax_op = partial(lax.dot, precision=lax.Precision.HIGHEST)
     self._CheckAgainstNumpy(lax_op, lax_reference.dot, args_maker, tol=tol)
@@ -680,7 +682,7 @@ def testDotAgainstNumpy(self, lhs_shape, rhs_shape, dtype, rng_factory):
           [(1, 2, 2, 3), (1, 2, 3, 1), [1], [1]],
           [(3, 2), (2, 4), [1], [0]],
       ]
-      for dtype in default_dtypes
+      for dtype in all_dtypes
       for rng_factory in [jtu.rand_small]))
   def testDotGeneralContractOnly(self, lhs_shape, rhs_shape, dtype,
                                  lhs_contracting, rhs_contracting, rng_factory):
@@ -705,7 +707,7 @@ def fun(lhs, rhs):
           ((3, 3, 2), (3, 2, 4), (([2], [1]), ([0], [0]))),
           ((3, 4, 2, 4), (3, 4, 3, 2), (([2], [3]), ([0, 1], [0, 1]))),
       ]
-      for dtype in default_dtypes
+      for dtype in all_dtypes
       for rng_factory in [jtu.rand_small]))
   def testDotGeneralContractAndBatch(self, lhs_shape, rhs_shape, dtype,
                                      dimension_numbers, rng_factory):
@@ -729,7 +731,7 @@ def fun(lhs, rhs):
           ((3, 3, 2), (3, 2, 4), (([2], [1]), ([0], [0]))),
           ((3, 4, 2, 4), (3, 4, 3, 2), (([2], [3]), ([0, 1], [0, 1]))),
       ]
-      for dtype in default_dtypes
+      for dtype in all_dtypes
       for rng_factory in [jtu.rand_small]))
   def testDotGeneralAgainstNumpy(self, lhs_shape, rhs_shape, dtype,
                                  dimension_numbers, rng_factory):