Avoid usage of deprecated jax.core APIs.

These APIs are deprecated as of JAX v0.10.0, replaced by equivalents in `jax.extend.core` (see https://docs.jax.dev/en/latest/jax.extend.html for details).

PiperOrigin-RevId: 900270194

Author: Jake VanderPlas

kfac_jax/_src/layers_and_loss_tags.py

@@ -19,6 +19,15 @@
 import jax
 import jax.extend as jex
 
+try:
+  # JAX v0.10.0 and newer
+  Effects: type[Any] = jex.core.Effects
+  no_effects: Effects = jex.core.no_effects
+except AttributeError:
+  # JAX v0.9.2 and older
+  Effects = jax.core.Effects
+  no_effects = jax.core.no_effects
+
 
 # Types for annotation
 T = TypeVar("T")
@@ -135,9 +144,9 @@ def abstract_eval(
       self,
       *args: Array,
       **params: Any,
-  ) -> tuple[Arrays, jax.core.Effects]:
+  ) -> tuple[Arrays, Effects]:
 
-    return get_loss_outputs(args, params), jax.core.no_effects
+    return get_loss_outputs(args, params), no_effects
 
   def _mlir_lowering(
       self,
@@ -336,10 +345,10 @@ def abstract_eval(
       self,
       *args: Array,
       **params: Any,
-  ) -> tuple[Array, jax.core.Effects]:
+  ) -> tuple[Array, Effects]:
     # For now we support only single output
     [output] = self.layer_data(args, params).outputs
-    return output, jax.core.no_effects
+    return output, no_effects
 
   def _batching(
       self,

kfac_jax/_src/tag_graph_matcher.py

@@ -34,10 +34,15 @@
     jax.__version_info__ if hasattr(jax, "__version_info__")
     else tuple(map(int, jax.__version__.split("."))))
 
-if jax_version >= (0, 5, 1):
+if jax_version >= (0, 10, 0):
+  DebugInfo = jex.core.DebugInfo
+  DropVar = jex.core.DropVar
+elif jax_version >= (0, 5, 1):
   DebugInfo = jax.core.DebugInfo
+  DropVar = jax.core.DropVar
 else:
   DebugInfo = jax.core.JaxprDebugInfo  #  pytype: disable=module-attr
+  DropVar = jax.core.DropVar
 
 
 HIGHER_ORDER_NAMES = ("cond", "while", "scan", "pjit", "xla_call", "xla_pmap")
@@ -376,7 +381,7 @@ def make_jax_graph(
           new_out_vars = []
           for v in eqn.outvars:
 
-            if isinstance(v, jax.core.DropVar):
+            if isinstance(v, DropVar):
               new_out_vars.append(make_var_func(v.aval))
             else:
               new_out_vars.append(v)
@@ -899,7 +904,7 @@ def read_env(
     if isinstance(v, jex.core.Literal):
       # Literals are values baked into the Jaxpr
       result.append(v.val)
-    elif isinstance(v, jax.core.DropVar):
+    elif isinstance(v, DropVar):
       result.append(None)
     else:
       result.append(env[v])

tests/test_graph_matcher.py

@@ -27,15 +27,21 @@
 Array = kfac_jax.utils.Array
 Shape = kfac_jax.utils.Shape
 
+try:
+  # JAX v0.10.0 and newer
+  DropVar = jex.core.DropVar
+except AttributeError:
+  # JAX v0.9.2 and older
+  DropVar = jax.core.DropVar
+
 
 class TestGraphMatcher(parameterized.TestCase):
   """Test class for the functions in `tag_graph_matcher.py`."""
 
   def check_equation_match(self, eqn1, vars_to_vars, vars_to_eqn):
     """Checks that equation is matched in the other graph."""
 
-    eqn1_out_vars = [v for v in eqn1.outvars
-                     if not isinstance(v, jax.core.DropVar)]
+    eqn1_out_vars = [v for v in eqn1.outvars if not isinstance(v, DropVar)]
     eqn2_out_vars = [vars_to_vars[v] for v in eqn1_out_vars]
     eqns = [vars_to_eqn[v] for v in eqn2_out_vars]
     self.assertTrue(all(e == eqns[0] for e in eqns[1:]))
@@ -124,8 +130,8 @@ def check_jaxpr_equal(self, jaxpr_1, jaxpr_2, map_output_vars: bool):
       for eqn1, eqn2 in zip(l1_eqns, l2_eqns):
         self.assertEqual(len(eqn1.outvars), len(eqn2.outvars))
         for v1, v2 in zip(eqn1.outvars, eqn2.outvars):
-          if isinstance(v1, jax.core.DropVar):
-            self.assertIsInstance(v2, jax.core.DropVar)
+          if isinstance(v1, DropVar):
+            self.assertIsInstance(v2, DropVar)
           elif isinstance(v1, jex.core.Literal):
             self.assertIsInstance(v2, jex.core.Literal)
             self.assertEqual(v1.aval, v2.aval)