Repository: google/jax
Issue: Type error in reverse-mode AD of `lax.dot_general` with `preferred_element_type`

Colab:
https://colab.research.google.com/gist/romanngg/64bae6efa7d4d364c68cd341f1be8c85/reverse_mode_dot_general_fail.ipynb

```python

import jax, jax.numpy as np

x = np.ones((), np.float16)

def f(x):
  return jax.lax.dot_general(x, x, (((), ()), ((), ())), preferred_element_type=np.float32)

jax.jacrev(f)(x)
```
causes
```
TypeError: lax.dot_general requires arguments to have the same dtypes, got float32, float16.
```
whereas forward pass and forward-mode AD work.

Changed files:
jax/_src/lax/convolution.py, jax/_src/lax/lax.py, tests/lax_autodiff_test.py, tests/lax_test.py

diff --git a/jax/_src/lax/convolution.py b/jax/_src/lax/convolution.py
index 301e16819..208baf131 100644
--- a/jax/_src/lax/convolution.py
+++ b/jax/_src/lax/convolution.py
@@ -417,11 +417,11 @@ def _conv_general_dilated_shape_rule(
 def _conv_general_dilated_dtype_rule(
     lhs, rhs, *, window_strides, padding, lhs_dilation, rhs_dilation,
     dimension_numbers, preferred_element_type, **unused_kwargs):
-  input_dtype = lax.naryop_dtype_rule(lax._input_dtype, [lax._any, lax._any],
-                                      'conv_general_dilated', lhs, rhs)
+  result_dtype = lax.naryop_dtype_rule(lax._input_dtype, [lax._any, lax._any],
+                                       'conv_general_dilated', lhs, rhs)
   if preferred_element_type is None:
-    return input_dtype
-  lax._validate_preferred_element_type(input_dtype, preferred_element_type)
+    return result_dtype
+  lax._validate_preferred_element_type(result_dtype, preferred_element_type)
   return preferred_element_type
 
 _conv_spec_transpose = lambda spec: (spec[1], spec[0]) + spec[2:]
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index d30b901c2..056446442 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -2399,23 +2399,30 @@ mlir.register_lowering(bitcast_convert_type_p, _bitcast_convert_type_lower)
 
 
 def _validate_preferred_element_type(input_dtype, preferred_element_type):
-
-  if dtypes.issubdtype(input_dtype, np.integer) and dtypes.issubdtype(preferred_element_type, np.floating):
+  if (dtypes.issubdtype(input_dtype, np.integer) and
+      dtypes.issubdtype(preferred_element_type, np.floating)):
     # Special-case integer->float multiply. This is allowed, and also allows
     # different signedness between input and output.
     pass
   else:
     allowed_types = (np.integer, np.floating, np.complexfloating)
-    if any(dtypes.issubdtype(input_dtype, t) and not dtypes.issubdtype(preferred_element_type, t) for t in allowed_types):
-      raise TypeError("Input type is incompatible with `preferred_element_type`. The compatible combinations of "
-                      "(input_type, preferred_element_type) are (integral, integral), (integral, floating), "
+    if any(dtypes.issubdtype(input_dtype, t) and not
+           dtypes.issubdtype(preferred_element_type, t) for t in allowed_types):
+      raise TypeError("Input type is incompatible with "
+                      "`preferred_element_type`. The compatible combinations "
+                      "of (input_type, preferred_element_type) are "
+                      "(integral, integral), (integral, floating), "
                       "(floating, floating), (complex, complex.")
-    if dtypes.issubdtype(input_dtype, np.signedinteger) and not dtypes.issubdtype(preferred_element_type, np.signedinteger):
-      raise TypeError("`preferred_element_type` must have the same signedness as the original type.")
+    if (dtypes.issubdtype(input_dtype, np.signedinteger) and
+        not dtypes.issubdtype(preferred_element_type, np.signedinteger)):
+      raise TypeError("`preferred_element_type` must have the same signedness "
+                      "as the original type.")
   input_bitwidth = np.dtype(input_dtype).itemsize
   preferred_bitwidth = np.dtype(preferred_element_type).itemsize
   if preferred_bitwidth < input_bitwidth:
-    raise TypeError("`preferred_element_type` must not be narrower than the original type.")
+    raise TypeError("`preferred_element_type` must not be narrower than the "
+                    "original type.")
+
 
 def _precision_config(precision):
   if precision is not None:
@@ -2506,10 +2513,46 @@ def tuple_delete(tup, idx):
 
 def _dot_general_dtype_rule(lhs, rhs, *, dimension_numbers, precision,
                             preferred_element_type: Optional[DTypeLike]):
-  input_dtype = naryop_dtype_rule(_input_dtype, [_any, _any], 'dot_general', lhs, rhs)
-  if preferred_element_type is None:
-    return input_dtype
-  _validate_preferred_element_type(input_dtype, preferred_element_type)
+  # We're mostly matching XLA's logic here, namely in shape_inference.cc and
+  # primitive_util.h's HigherPrecisionType, e.g.
+  # https://github.com/openxla/xla/blob/ea3a841768d0dcf192e5820c9b25c34c73f2226a/xla/primitive_util.h#L329
+  def type_properties(dt):
+    c = _real_dtype(dt) if dtypes.issubdtype(dt, np.complexfloating) else dt
+    return (dtypes.issubdtype(dt, np.complexfloating),
+            dtypes.finfo(c).maxexp if dtypes.issubdtype(c, np.floating) else -1,
+            dtypes.finfo(c).nmant  if dtypes.issubdtype(c, np.floating) else -1,
+            _bit_width(c),
+            not dtypes.issubdtype(c, np.unsignedinteger))
+  lhs_prop, rhs_prop = type_properties(lhs.dtype), type_properties(rhs.dtype)
+  if lhs_prop > rhs_prop:
+    result_dtype = lhs.dtype
+  elif rhs_prop > lhs_prop:
+    result_dtype = rhs.dtype
+  else:
+    if lhs.dtype != rhs.dtype:
+      raise TypeError(
+          f"lax.dot_general argument type error: {lhs.dtype}, {rhs.dtype}")
+    result_dtype = lhs.dtype
+
+  return _maybe_upcast(result_dtype, preferred_element_type)
+
+def _bit_width(d):
+  if dtypes.issubdtype(d, np.inexact): return dtypes.finfo(d).bits
+  elif dtypes.issubdtype(d, np.integer): return dtypes.iinfo(d).bits
+  elif d == np.dtype('bool'): return 1
+  else: assert False, d  # should be unreachable, open an issue!
+
+def _maybe_upcast(result_dtype, preferred_element_type):
+  # replicates the logic in shape_inference.cc's MaybeUpcast
+  if (preferred_element_type is None or
+      result_dtype == preferred_element_type):
+    return result_dtype
+  if (not dtypes.issubdtype(result_dtype, np.floating) and
+      _bit_width(preferred_element_type) < _bit_width(result_dtype)):
+    raise TypeError("`preferred_element_type` must not be narrower than the "
+                    "original type, got preferred_element_type of "
+                    f"{preferred_element_type} for result type of "
+                    f"{result_dtype}.")
   return preferred_element_type
 
 def _dot_general_transpose_lhs(g, x, y, *, dimension_numbers, precision,
@@ -2526,19 +2569,24 @@ def _dot_general_transpose_lhs(g, x, y, *, dimension_numbers, precision,
   dims = ((ans_y, y_kept), (ans_batch, y_batch))
   x_contract_sorted_by_y = list(np.take(x_contract, np.argsort(y_contract)))  # type: ignore[arg-type]
   out_axes = np.argsort(list(x_batch) + x_kept + x_contract_sorted_by_y)
-  return transpose(dot_general(g, y, dims, precision=precision,
-                               preferred_element_type=preferred_element_type),
-                   tuple(out_axes))
+  x_bar = transpose(dot_general(g, y, dims, precision=precision,
+                                preferred_element_type=preferred_element_type),
+                    tuple(out_axes))
+  if x_bar.dtype != x.aval.dtype:
+    x_bar = _convert_element_type(x_bar, x.aval.dtype, x.aval.weak_type)
+  return x_bar
 
 def _dot_general_transpose_rhs(g, x, y, *, dimension_numbers, precision,
                                preferred_element_type: Optional[DTypeLike]):
   (x_contract, y_contract), (x_batch, y_batch) = dimension_numbers
   swapped_dimension_numbers = ((y_contract, x_contract), (y_batch, x_batch))
-  return _dot_general_transpose_lhs(
+  y_bar = _dot_general_transpose_lhs(
     g, y, x, dimension_numbers=swapped_dimension_numbers, precision=precision,
     preferred_element_type=preferred_element_type,
     swap_ans=True)
-
+  if y_bar.dtype != y.aval.dtype:
+    y_bar = _convert_element_type(y_bar, y.aval.dtype, y.aval.weak_type)
+  return y_bar
 
 def _dot_general_batch_rule(batched_args, batch_dims, *, dimension_numbers,
                             precision,
@@ -2667,19 +2715,47 @@ def _dot_general_lower(ctx, lhs, rhs, *, dimension_numbers,
                        precision, preferred_element_type: Optional[np.dtype]):
   del preferred_element_type  # Implied by the output aval
   lhs_aval, rhs_aval = ctx.avals_in
+  lhs_dtype, rhs_dtype = lhs_aval.dtype, rhs_aval.dtype
   aval_out, = ctx.avals_out
   (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers
 
+  # TODO(b/...): JAX's dot_general primitive accepts the same input dtype
+  # combinations that are accepted in XLA's shape_inference.cc (the canonical
+  # reference for the HLO type system), but actually different XLA platforms
+  # fail on codegen for different accepted cases. To handle those cases, we
+  # insert ConvertOps on the input, in a platform-dependent way.
+  if lhs_dtype != rhs_dtype:
+    if ctx.module_context.platform == "tpu":
+      handled = lambda dt: (dtypes.issubdtype(dt, np.floating) or
+                            dtypes.issubdtype(dt, np.integer))
+      if not (handled(lhs_dtype) and handled(rhs_dtype)):
+        dt = mlir.dtype_to_ir_type(aval_out.dtype)
+        lhs = hlo.ConvertOp(ir.RankedTensorType.get(lhs_aval.shape, dt), lhs
+                            ).result
+        rhs = hlo.ConvertOp(ir.RankedTensorType.get(rhs_aval.shape, dt), rhs
+                            ).result
+        lhs_dtype = rhs_dtype = aval_out.dtype
+    else:  # cpu and gpu
+      dt = mlir.dtype_to_ir_type(aval_out.dtype)
+      lhs = hlo.ConvertOp(ir.RankedTensorType.get(lhs_aval.shape, dt), lhs
+                          ).result
+      rhs = hlo.ConvertOp(ir.RankedTensorType.get(rhs_aval.shape, dt), rhs
+                          ).result
+      lhs_dtype = rhs_dtype = aval_out.dtype
+
   # TODO(b/195364460): Work around slow XLA/CPU implementation of float16 matmul
   if ctx.module_context.platform == "cpu":
-    if lhs_aval.dtype == np.float16:
+    if lhs_dtype == np.float16:
       f32 = mlir.dtype_to_ir_type(np.dtype(np.float32))
       lhs = hlo.ConvertOp(ir.RankedTensorType.get(lhs_aval.shape, f32),
                           lhs).result
-    if rhs_aval.dtype == np.float16:
+      lhs_dtype = np.dtype('float32')
+    if rhs_dtype == np.float16:
       f32 = mlir.dtype_to_ir_type(np.dtype(np.float32))
       rhs = hlo.ConvertOp(ir.RankedTensorType.get(rhs_aval.shape, f32),
                           rhs).result
+      rhs_dtype = np.dtype('float32')
+
   dot_dnums = hlo.DotDimensionNumbers.get(
       lhs_batching_dimensions=list(lhs_batch),
       rhs_batching_dimensions=list(rhs_batch),
diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
index 52935a99a..46f8d324d 100644
--- a/tests/lax_autodiff_test.py
+++ b/tests/lax_autodiff_test.py
@@ -424,6 +424,14 @@ class LaxAutodiffTest(jtu.JaxTestCase):
     s = str(jax.make_jaxpr(pullback)(gresult))
     assert "Precision.HIGHEST" in s
 
+  def testDotPreferredElementType(self):
+    # https://github.com/google/jax/issues/10818
+    x = jax.numpy.ones((), jax.numpy.float16)
+    def f(x):
+      return jax.lax.dot_general(x, x, (((), ()), ((), ())),
+                                 preferred_element_type=jax.numpy.float32)
+    jax.jacrev(f)(x)  # don't crash!
+
   @jtu.sample_product(
     shape=[(), (2, 3)],
     dtype=float_dtypes,
diff --git a/tests/lax_test.py b/tests/lax_test.py
index f8e0f51c9..d380a687a 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -976,12 +976,17 @@ class LaxTest(jtu.JaxTestCase):
                                            padding=(3, 3))
 
   @jtu.sample_product(
-      [
-          dict(lhs_shape=lhs_shape, rhs_shape=rhs_shape)
-          for lhs_shape in [(3,), (4, 3)]
-          for rhs_shape in [(3,), (3, 6)]
-      ],
-      dtype=lax_test_util.all_dtypes,
+      [dict(lhs_shape=lhs_shape, rhs_shape=rhs_shape)
+       for lhs_shape in [(3,), (4, 3)] for rhs_shape in [(3,), (3, 6)]],
+      [dict(lhs_dtype=lhs_dtype, rhs_dtype=rhs_dtype)
+       for lhs_dtype, rhs_dtype in
+       itertools.chain(
+           itertools.product(lax_test_util.int_dtypes +
+                             lax_test_util.float_dtypes +
+                             lax_test_util.complex_dtypes +
+                             lax_test_util.uint_dtypes,
+                             repeat=2),
+           zip(lax_test_util.bool_dtypes, lax_test_util.bool_dtypes))],
       precision=[
           None,
           lax.Precision.DEFAULT,
@@ -990,9 +995,9 @@ class LaxTest(jtu.JaxTestCase):
           (lax.Precision.DEFAULT, lax.Precision.HIGHEST),
       ],
   )
-  def testDot(self, lhs_shape, rhs_shape, dtype, precision):
+  def testDot(self, lhs_shape, rhs_shape, lhs_dtype, rhs_dtype, precision):
     rng = jtu.rand_default(self.rng())
-    args_maker = lambda: [rng(lhs_shape, dtype), rng(rhs_shape, dtype)]
+    args_maker = lambda: [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)]
     self._CompileAndCheck(partial(lax.dot, precision=precision), args_maker)
 
   @jtu.sample_product(
@@ -1001,7 +1006,8 @@ class LaxTest(jtu.JaxTestCase):
     [dict(dtype=d, preferred_element_type=p)
      for d, p in preferred_type_combinations],
   )
-  def testDotPreferredElement(self, lhs_shape, rhs_shape, dtype, preferred_element_type):
+  def testDotPreferredElement(self, lhs_shape, rhs_shape, dtype,
+                              preferred_element_type):
     if (not config.x64_enabled and
        (dtype == np.float64 or preferred_element_type == np.float64
         or dtype == np.int64 or preferred_element_type == np.int64)):
@@ -1011,7 +1017,8 @@ class LaxTest(jtu.JaxTestCase):
       raise SkipTest("np.complex128 is not yet supported on TPU")
     if jtu.device_under_test() == "gpu":
       # TODO(b/189287598)
-      raise SkipTest("dot_general with preferred_element_type returns NaN non-deterministically on GPU")
+      raise SkipTest("dot_general with preferred_element_type returns NaN "
+                     "non-deterministically on GPU")
     rng = jtu.rand_default(self.rng())
     x = rng(lhs_shape, dtype)
     y = rng(rhs_shape, dtype)