distrax: migrate from deprecated jax.linear_util to jax.extend.linear…

…_util

PiperOrigin-RevId: 577275354

distrax/_src/utils/transformations.py

@@ -53,6 +53,12 @@
 import jax
 import jax.numpy as jnp
 
+try:
+  # jax >= 0.4.16
+  from jax.extend import linear_util as lu  # pylint: disable=g-import-not-at-top
+except ImportError:
+  from jax import linear_util as lu  # pylint: disable=g-import-not-at-top
+
 
 _inverse_registry = {
     # unary ops
@@ -258,7 +264,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 = [lu.wrap_init(
           functools.partial(_interpret_inverse, call_jaxpr, ()))]
       prim_inv = eqn.primitive