Replace deprecated jax.linear_util.wrap_init with jax.extend.linear_u…

…til.wrap_init. PiperOrigin-RevId: 562017973
google-deepmind · Sep 1, 2023 · 175cace · 175cace
1 parent 93c54a8
commit 175cace
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/distrax/_src/utils/transformations.py b/distrax/_src/utils/transformations.py
@@ -258,7 +258,7 @@ def write(var, val):
     # if primitive is an xla_call, get subexpressions and evaluate recursively
     call_jaxpr, params = _extract_call_jaxpr(eqn.primitive, params)
     if call_jaxpr:
-      subfuns = [jax.linear_util.wrap_init(
+      subfuns = [jax.extend.linear_util.wrap_init(
           functools.partial(_interpret_inverse, call_jaxpr, ()))]
       prim_inv = eqn.primitive