* Introducing a more general KroneckerFactored block class.

PiperOrigin-RevId: 531184538
google-deepmind · May 16, 2023 · 3e6b4e7 · 3e6b4e7
1 parent 95df497
commit 3e6b4e7
Show file tree

Hide file tree

Showing 5 changed files with 313 additions and 253 deletions.
diff --git a/kfac_jax/__init__.py b/kfac_jax/__init__.py
@@ -69,6 +69,7 @@
 ScaledIdentity = curvature_blocks.ScaledIdentity
 Diagonal = curvature_blocks.Diagonal
 Full = curvature_blocks.Full
+KroneckerFactored = curvature_blocks.KroneckerFactored
 TwoKroneckerFactored = curvature_blocks.TwoKroneckerFactored
 NaiveDiagonal = curvature_blocks.NaiveDiagonal
 NaiveFull = curvature_blocks.NaiveFull
@@ -146,6 +147,7 @@
     "ScaledIdentity",
     "Diagonal",
     "Full",
+    "KroneckerFactored",
     "TwoKroneckerFactored",
     "NaiveDiagonal",
     "NaiveFull",