Repository: google/jax

Issue:
jnp.float16 hangs matmul
I'm not sure if this is a bug or expected behavior but couldn't find a similar issue here so thought I'd post

Repro

The below code just runs seemingly forever where the `float32` and `bfloat16` variants run just fine. Is `float16` supported? I tried both a GPU and TPU in collab and noticed the same issue.

```python
import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random

import jax; jax.config.update('jax_platform_name', 'cpu')

key = random.PRNGKey(0)
x = random.normal(key, (10,))
print(x)

size = 3000
x = random.normal(key, (size, size), dtype=jnp.float16)
%timeit jnp.dot(x, x.T).block_until_ready()  # runs on the GPU
```

`x = random.normal(key, (size, size), dtype=jnp.float32)` runs just fine




Discussion:
Thanks for the report & the clear repro. This seems like something that @hawkinsp might have ideas about.

Resolving patch:
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
@@ -3552,6 +3552,23 @@ def _dot_general_translation_rule(c, lhs, rhs, *, dimension_numbers, precision,
                          precision_config=_precision_config(precision),
                          preferred_element_type=preferred_element_type)
 
+def _dot_general_cpu_translation_rule(c, lhs, rhs, *, dimension_numbers, precision,
+                                  preferred_element_type: Optional[DType]):
+  if preferred_element_type is not None:
+    preferred_element_type = xla_client.dtype_to_etype(preferred_element_type)
+
+  # TODO(b/195364460): Work around slow XLA/CPU implementation of float16 matmul
+  if c.get_shape(lhs).numpy_dtype() == np.float16:
+    lhs = xops.ConvertElementType(lhs, xla_client.dtype_to_etype(np.float32))
+    rhs = xops.ConvertElementType(rhs, xla_client.dtype_to_etype(np.float32))
+    preferred_element_type = (preferred_element_type or
+                              xla_client.dtype_to_etype(np.float16))
+
+  return xops.DotGeneral(lhs, rhs,
+                         xc.make_dot_dimension_numbers(dimension_numbers),
+                         precision_config=_precision_config(precision),
+                         preferred_element_type=preferred_element_type)
+
 def _dot_general_masking_rule(padded_vals, logical_shapes, *, dimension_numbers,
                               precision,
                               preferred_element_type: Optional[DType]):
@@ -3572,6 +3589,9 @@ def _dot_general_masking_rule(padded_vals, logical_shapes, *, dimension_numbers,
                _dot_general_transpose_lhs, _dot_general_transpose_rhs)
 batching.primitive_batchers[dot_general_p] = _dot_general_batch_rule
 masking.masking_rules[dot_general_p] = _dot_general_masking_rule
+xla.backend_specific_translations["cpu"][dot_general_p] = \
+  _dot_general_cpu_translation_rule
+
 
 def _broadcast_in_dim_shape_rule(operand, *, shape, broadcast_dimensions):
   _check_shapelike('broadcast_in_dim', 'shape', shape)


Test patch:
diff --git a/jax/experimental/jax2tf/tests/jax2tf_limitations.py b/jax/experimental/jax2tf/tests/jax2tf_limitations.py
--- a/jax/experimental/jax2tf/tests/jax2tf_limitations.py
+++ b/jax/experimental/jax2tf/tests/jax2tf_limitations.py
@@ -405,7 +405,11 @@ def dot_general(cls, harness: primitive_harness.Harness):
             devices="gpu",
             modes=("eager", "graph", "compiled"),
             enabled=(harness.params["preferred_element_type"] is not None),
-            skip_comparison=True)
+            skip_comparison=True),
+        # JAX performs float16 matmuls in float32 on CPU, so the JAX result
+        # may be more precise.
+        custom_numeric(dtypes=[np.float16], devices=["cpu"], tol=1e-2,
+                       modes=("eager", "graph", "compiled")),
     ]
 
   @classmethod