Reverts 73962b7

PiperOrigin-RevId: 703100851

Author: chr1sj0nes

jax/_src/lax/lax.py

@@ -923,8 +923,9 @@ def accumulation_type(self) -> DTypeLike | None:
       case _:
         return np.float32
 
-  @property
-  def supported_output_types(self) -> tuple[DTypeLike, ...] | None:
+  def supported_output_types(
+      self, lhs_dtype: DTypeLike, rhs_dtype: DTypeLike
+  ) -> tuple[DTypeLike, ...] | None:
     match self:
       case (
           DotAlgorithmPreset.ANY_F8_ANY_F8_F32
@@ -941,7 +942,17 @@ def supported_output_types(self) -> tuple[DTypeLike, ...] | None:
             dtypes.float8_e4m3b11fnuz,
         )
       case DotAlgorithmPreset.F16_F16_F32:
-        return (np.float32, np.float16)
+        # F16 output is only supported with F16 inputs.
+        if dtypes.promote_types(lhs_dtype, rhs_dtype) == np.float16:
+          return (np.float32, np.float16)
+        else:
+          return (np.float32,)
+      case DotAlgorithmPreset.BF16_BF16_F32:
+        # BF16 output is only supported with BF16 inputs.
+        if dtypes.promote_types(lhs_dtype, rhs_dtype) == dtypes.bfloat16:
+          return (np.float32, dtypes.bfloat16)
+        else:
+          return (np.float32,)
       case _:
         accumulation_type = self.accumulation_type
         return None if accumulation_type is None else (accumulation_type,)
@@ -3713,25 +3724,19 @@ def get_algorithm_compute_types(
         algorithm.accumulation_type,
     )
 
-  supported_output_types = algorithm.supported_output_types
-
-  if algorithm == DotAlgorithmPreset.BF16_BF16_F32:
-    # If dtype is anything other than float32, it will be cast to bfloat16.
-    if np.dtype(lhs_dtype) != np.float32:
-      supported_output_types = (np.float32, dtypes.bfloat16)
-
   def maybe_convert_dtype(input_dtype, target_dtypes):
     if target_dtypes is None:
       return input_dtype
     if np.dtype(input_dtype) in map(np.dtype, target_dtypes):
       return input_dtype
     return target_dtypes[0]
 
-  return (
-      maybe_convert_dtype(lhs_dtype, algorithm.supported_lhs_types),
-      maybe_convert_dtype(rhs_dtype, algorithm.supported_rhs_types),
-      maybe_convert_dtype(out_dtype, supported_output_types),
+  lhs_dtype = maybe_convert_dtype(lhs_dtype, algorithm.supported_lhs_types)
+  rhs_dtype = maybe_convert_dtype(rhs_dtype, algorithm.supported_rhs_types)
+  out_type = maybe_convert_dtype(
+      out_dtype, algorithm.supported_output_types(lhs_dtype, rhs_dtype)
   )
+  return lhs_dtype, rhs_dtype, out_type
 
 
 def _dot_general_lower(ctx, lhs, rhs, *, dimension_numbers,