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.


Discussion:
+1, any updates? 

``` 
x = jnp.einsum(
        "blhd,bmd->bhlm",
        q,
        k,
        preferred_element_type=jnp.float32
    )
````

where q and k are bf16 crashes on backward pass with the 
```TypeError: lax.dot_general requires arguments to have the same dtypes, got float32, bfloat16.```
you can work around this if you do q.astype(jnp.float32), k.astype(jnp.float32) before the einsum
Thanks for the ping!

This should be easy to fix once we figure out the logic to implement. But I'm unsure of the logic we want. Maybe you can tell me!

Say we transpose (essentially, reverse-mode differentiate) a dot of the form
```
z:f32[3,5] = dot_general [preferred_element_type=f32] x:bf16[3,4] y:bf16[4,5]
```
with respect to its second argument, where the contraction is happening over the axes of the same size. That means we need to take a `z_bar:f32[3,5]` and an `x:bf16[3,4]` and ultimately produce a `y_bar:bf16[4,5]`.

At one point, probably in the ancient past, I think HLO didn't accept `dot_general` applications with different element types. Let's say that constraint is not present. Here are two different things we could transpose the above to:

```
# Option 1: change the preferred_element_type, so accumulation (and output) is now in bf16
y_bar:bf16[4,5] = dot_general [preferred_element_type=bf16] x:bf16[3,4] z_bar:f32[3,5]

# Option 2: keep the preferred_element_type (so accum is in f32), do a cast afterward
y_bar_temp:f32[4,5] = dot_general [preferred_element_type=f32] x:bf16[3,4] z_bar:f32[3,5]
y_bar:bf16[4,5] = convert_element_type[new_dtype=bf16] y_bar_temp:f32[4,5]
```

Is it clear to you which one of those is better? Is it ambiguous?
I think `preferred_element_type` doesn't always mean that accumulation happens in that precision (e.g. accumulation of bf16 matmuls on TPU is always in f32)
But I asked an LLM and that's what it told me!

Just kidding, I was going off of our docs [here](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.conv_with_general_padding.html):
> preferred_element_type (Optional[Any]) – Optional. Either None, which means the default accumulation type for the input types, or a datatype, indicating to accumulate results to and return a result with that datatype.

Maybe our docs are wrong (EDIT: or I'm misunderstanding them).
It should be just as fast to cast the inputs to f32, right?   Wouldn't XLA do some magic optimization to avoid materializing the f32 input in HBM?
That seems likely. But I'm not familiar enough (with what XLA does, on every backend) to know the best JAX-level policy (or if it even matters given XLA magic).

Let me ask the wizard Blake...
Blake and his fellow wizard @hawkinsp advised! TL;DR we will try Option 2 above.
I'll make a fix for this today, unless it's more urgent and you need it ASAP.
it is not super urgent. thank you for your help!

Resolving patch:
diff --git a/jax/_src/lax/convolution.py b/jax/_src/lax/convolution.py
--- 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
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -2399,23 +2399,30 @@ def _bitcast_convert_type_lower(ctx, operand, *, new_dtype):
 
 
 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),


Test patch:
diff --git a/tests/lax_autodiff_test.py b/tests/lax_autodiff_test.py
--- a/tests/lax_autodiff_test.py
+++ b/tests/lax_autodiff_test.py
@@ -424,6 +424,14 @@ def testDotGeneralContractAndBatchGrads(self, lhs_shape, rhs_shape, dtype,
     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
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -976,12 +976,17 @@ def testConvInvalidPadding(self):
                                            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 @@ def testConvInvalidPadding(self):
           (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 @@ def testDot(self, lhs_shape, rhs_shape, dtype, precision):
     [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 @@ def testDotPreferredElement(self, lhs_shape, rhs_shape, dtype, preferred_element
       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)