Direct quantization for FP8 Dense Layer.

google · Sep 3, 2024 · 6f5cee1 · 6f5cee1
1 parent 239d4e6
commit 6f5cee1
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,