Direct quantization for FP8 Dense Layer.

google · Aug 30, 2024 · 2b7ff5a · 2b7ff5a
1 parent 239d4e6
commit 2b7ff5a
Show file tree

Hide file tree

Showing 3 changed files with 341 additions and 45 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,