Merge pull request #3922 from wenscarl:direct_quant

PiperOrigin-RevId: 671087442
google · Sep 4, 2024 · c44b916 · c44b916
2 parents 35c4edf + e9d19ac
commit c44b916
Show file tree

Hide file tree

Showing 3 changed files with 333 additions and 44 deletions.
diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py
@@ -86,8 +86,11 @@
 )
 from .batch_apply import BatchApply as BatchApply
 from .combinators import Sequential as Sequential
-from .fp8_ops import Fp8DotGeneralOp as Fp8DotGeneralOp
-from .fp8_ops import NANOOFp8DotGeneralOp as NANOOFp8DotGeneralOp
+from .fp8_ops import (
+    Fp8DotGeneralOp as Fp8DotGeneralOp,
+    Fp8DirectDotGeneralOp as Fp8DirectDotGeneralOp,
+    NANOOFp8DotGeneralOp as NANOOFp8DotGeneralOp,
+)
 from .initializers import (
     ones_init as ones_init,
     ones as ones,