Add missing decorator

google · Nov 7, 2023 · aa9427a · aa9427a
1 parent 786a111
commit aa9427a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flax/linen/fp8_ops.py b/flax/linen/fp8_ops.py
@@ -135,7 +135,7 @@ def dot_general_with_precision(
     lhs, rhs, dimension_numbers, precision=lax.Precision.DEFAULT
   )
 
-
+@dot_general_with_precision.defjvp
 def dot_general_with_precision_jvp(
   dimension_numbers, precision, preferred_element_type, primals, tangents
 ):