* Revamping the graph matching code to be able to detect layers and r…

…egister tag in arbitrary higher-order Jax primitives. * A few minor typos. * Bumping jax and jaxlib version to 0.3.15 PiperOrigin-RevId: 460947146
google-deepmind · Aug 11, 2022 · cc2cd86 · cc2cd86
1 parent 698fcc6
commit cc2cd86
Show file tree

Hide file tree

Showing 9 changed files with 1,219 additions and 541 deletions.
diff --git a/kfac_jax/_src/curvature_estimator.py b/kfac_jax/_src/curvature_estimator.py
@@ -983,7 +983,7 @@ def _finalize(self, func_args: utils.FuncArgs):
 
   def init(
       self,
-      rng: chex.Array,
+      rng: chex.PRNGKey,
       func_args: utils.FuncArgs,
       exact_powers_to_cache: Optional[curvature_blocks.ScalarOrSequence],
       approx_powers_to_cache: Optional[curvature_blocks.ScalarOrSequence],