Add graph scan support for arbitrary number of input and output dimen…

…sions for dense layers PiperOrigin-RevId: 697590463 Change-Id: Id419b8b3a7288651a142c889d6bd2964c5da97b7
google-deepmind · Dec 8, 2024 · 9a1deec · 9a1deec
1 parent 38df5df
commit 9a1deec
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/ferminet/curvature_tags_and_blocks.py b/ferminet/curvature_tags_and_blocks.py
@@ -38,7 +38,13 @@ def register_qmc(y, x, w, **kwargs):
   return kfac_jax.register_dense(y, x, w, variant="qmc", **kwargs)
 
 
-_dense = kfac_jax.tag_graph_matcher._dense  # pylint: disable=protected-access
+_dense = functools.partial(
+    kfac_jax.tag_graph_matcher._dense,  # pylint: disable=protected-access
+    axes=1,
+    with_reshape=False,
+)
+
+
 _repeated_dense_parameter_extractor = functools.partial(
     kfac_jax.tag_graph_matcher._dense_parameter_extractor,  # pylint: disable=protected-access
     variant="repeated_dense",