Merge pull request #3960 from chiamp:fp8

PiperOrigin-RevId: 639953317
google · Jun 3, 2024 · 40c78bd · 40c78bd
2 parents 1619778 + 5e8e595
commit 40c78bd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flax/linen/fp8_ops.py b/flax/linen/fp8_ops.py
@@ -41,7 +41,7 @@ class Fp8MetaTyRules:
   def physical_element_aval(dtype) -> core.ShapedArray:
     return core.ShapedArray((), dtype.float_dtype)
 
-  if jax.__version_info__ < (0, 4, 28):
+  if jax.__version_info__ < (0, 4, 29):
     @staticmethod
     def replicate_trailing_dims(ctx, val, aval):
       del ctx, aval