[ROCm] Add float8_e4m3fnuz and float8_e5m2fnuz support for Rocm

Author: Ruturaj4

jax/_src/lax/lax.py

@@ -2913,8 +2913,9 @@ def _dot_general_lower(ctx, lhs, rhs, *, dimension_numbers,
                        precision, preferred_element_type: np.dtype | None,
                        platform: str = "default"):
   def _is_fp8_mixed_precision_matmul(_lhs_dtypes, _rhs_dtypes):
-      fp8_dtypes = (dtypes.float8_e4m3fn, dtypes.float8_e5m2)
-      return _lhs_dtypes in fp8_dtypes and _rhs_dtypes in fp8_dtypes
+    fp8_dtypes = (dtypes.float8_e4m3fn, dtypes.float8_e5m2,
+                  dtypes.float8_e5m2fnuz, dtypes.float8_e4m3fnuz)
+    return _lhs_dtypes in fp8_dtypes and _rhs_dtypes in fp8_dtypes
   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

tests/lax_test.py

@@ -1079,11 +1079,18 @@ def testDotPreferredElement(self, lhs_shape, rhs_shape, dtype,
      for lhs_shape in [(3,), (4, 3)] for rhs_shape in [(3,), (3, 6)]],
     [dict(dtype_lhs=dtype_lhs, dtype_rhs=dtype_rhs)
      for dtype_lhs, dtype_rhs in [(dtypes.float8_e4m3fn, dtypes.float8_e5m2),
-                                  (dtypes.float8_e5m2, dtypes.float8_e4m3fn)]],
+                                  (dtypes.float8_e5m2, dtypes.float8_e4m3fn),
+                                  (dtypes.float8_e4m3fnuz, dtypes.float8_e5m2fnuz),
+                                  (dtypes.float8_e5m2fnuz, dtypes.float8_e4m3fnuz)]],
   )
   def test_mixed_fp8_dot_general(self, lhs_shape, rhs_shape, dtype_lhs, dtype_rhs):
     if jtu.test_device_matches(["tpu"]):
         raise SkipTest("Mixed fp8 precision matmul is not yet supported on TPU")
+    if not jtu.is_device_rocm() and (
+        dtype_lhs in [dtypes.float8_e4m3fnuz, dtypes.float8_e5m2fnuz] or
+        dtype_rhs in [dtypes.float8_e4m3fnuz, dtypes.float8_e5m2fnuz]
+    ):
+        raise SkipTest("float8_e4m3fnuz and float8_e5m2fnuz types are only supported on ROCm")
     rng = jtu.rand_default(self.rng())
     lhs = rng(lhs_shape, dtype=dtype_lhs)
     rhs = rng(rhs_shape, dtype=dtype_rhs)