Merge pull request #3440 from chiamp:fix_fp8

PiperOrigin-RevId: 576953704
google · Oct 26, 2023 · 9064108 · 9064108
2 parents 738078c + e7e33e1
commit 9064108
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flax/linen/fp8_ops.py b/flax/linen/fp8_ops.py
@@ -174,13 +174,13 @@ def __call__(self, *args, **kwargs) -> jnp.ndarray:
     k_qdq = in_qdq(
         comp_dtype, k, self.kernel_scale.value, self.kernel_amax_history.value
     )
-    y_qdq = lax.dot_general(x_qdq, k_qdq, dimension_numbers, precision)
+    y_qdq = lax.dot_general(x_qdq, k_qdq, dimension_numbers, precision) # type: ignore
     y = out_qdq(
         comp_dtype,
         y_qdq,
         self.output_grad_scale.value,
         self.output_grad_amax_history.value
     )
 
-    return y
+    return y # type: ignore