jax.tree_util.* -> jax.tree.*

docs/binned_likelihood.md

@@ -59,7 +59,7 @@ def NLL(dynamic_params, static_params, hists, observation):
     # second product of Eq. 1 (constraint)
     constraints = evm.loss.get_log_probs(model)
     # for parameters with `.value.size > 1` (jnp.sum the constraints)
-    constraints = jtu.tree_map(jnp.sum, constraints)
+    constraints = jax.tree.map(jnp.sum, constraints)
     loss_val += evm.util.sum_over_leaves(constraints)
     return -jnp.sum(loss_val)
 ```

docs/building_blocks.md

@@ -101,7 +101,6 @@ Correlate a Parameter
     A more general case of correlating any PyTree of parameters is implemented as follows:
     ```{code-block} python
     from typing import NamedTuple
-    import jax.tree_util as jtu
 
 
     class Params(NamedTuple):
@@ -110,11 +109,11 @@ Correlate a Parameter
 
     params = Params(mu=evm.Parameter(1.0), syst=evm.NormalParameter(0.0))
 
-    flat_params, tree_def = jtu.tree_flatten(params, evm.parameter.is_parameter)
+    flat_params, tree_def = jax.tree.flatten(params, evm.parameter.is_parameter)
 
     # correlate the parameters
     correlated_flat_params = evm.parameter.correlate(*flat_params)
-    correlated_params = jtu.tree_unflatten(tree_def, correlated_flat_params)
+    correlated_params = jax.tree.unflatten(tree_def, correlated_flat_params)
 
     # now correlated_params.mu and correlated_params.syst are correlated,
     # they share the same value

examples/dnn_weights_constraint.py

@@ -36,7 +36,7 @@ def loss_fn(model, x, y):
     mse = jax.numpy.mean((y - pred_y) ** 2)
     constraints = evm.loss.get_log_probs(model)
     # sum them all up for each weight
-    constraints = jax.tree_util.tree_map(jnp.sum, constraints)
+    constraints = jax.tree.map(jnp.sum, constraints)
     return mse + evm.util.sum_over_leaves(constraints)
 
 

src/evermore/modifier.py

@@ -8,7 +8,6 @@
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-import jax.tree_util as jtu
 from jaxtyping import Array, ArrayLike, PyTree
 
 from evermore.custom_types import ModifierLike, OffsetAndScale
@@ -73,7 +72,6 @@ class ModifierBase(ApplyFn, MatMulCompose, AbstractModifier, SupportsTreescope):
 
         import equinox as eqx
         import jax.numpy as jnp
-        import jax.tree_util as jtu
         from jaxtyping import Array
 
         import evermore as evm
@@ -86,7 +84,7 @@ class Clip(evm.modifier.ModifierBase):
 
             def offset_and_scale(self, hist: Array) -> evm.custom_types.OffsetAndScale:
                 os = self.modifier.offset_and_scale(hist)
-                return jtu.tree_map(lambda x: jnp.clip(x, self.min_sf, self.max_sf), os)
+                return jax.tree.map(lambda x: jnp.clip(x, self.min_sf, self.max_sf), os)
 
 
         parameter = evm.Parameter(value=1.1)
@@ -201,7 +199,7 @@ def offset_and_scale(self, hist: Array) -> OffsetAndScale:
         def _where(true: Array, false: Array) -> Array:
             return jnp.where(self.condition, true, false)
 
-        return jtu.tree_map(_where, true_os, false_os)
+        return jax.tree.map(_where, true_os, false_os)
 
 
 class BooleanMask(ModifierBase):
@@ -280,7 +278,7 @@ class Transform(ModifierBase):
 
     def offset_and_scale(self, hist: Array) -> OffsetAndScale:
         os = self.modifier.offset_and_scale(hist)
-        return jtu.tree_map(self.transform_fn, os)
+        return jax.tree.map(self.transform_fn, os)
 
 
 class TransformOffset(ModifierBase):
@@ -377,7 +375,7 @@ def offset_and_scale(self, hist: Array) -> OffsetAndScale:
         groups = defaultdict(list)
         # first group modifiers into same tree structures
         for mod in self.unroll_modifiers():
-            groups[hash(jtu.tree_structure(mod))].append(mod)
+            groups[hash(jax.tree.structure(mod))].append(mod)
         # then do the `jax.lax.scan` loops
         for _, group_mods in groups.items():
             # skip empty groups

src/evermore/parameter.py

@@ -6,7 +6,6 @@
 import equinox as eqx
 import jax
 import jax.numpy as jnp
-import jax.tree_util as jtu
 from jaxtyping import Array, ArrayLike, PRNGKeyArray, PyTree
 
 from evermore.custom_types import PDFLike
@@ -139,15 +138,15 @@ def model(diffable, static, hists) -> Array:
             ...
     """
     # 1. set the filter_spec to False for all non-static leaves
-    filter_spec = jtu.tree_map(lambda _: False, tree)
+    filter_spec = jax.tree.map(lambda _: False, tree)
 
     # 2. set the filter_spec to True for each parameter value
     def _replace_value(leaf: Any) -> Any:
         if isinstance(leaf, Parameter):
             leaf = eqx.tree_at(lambda p: p.value, leaf, not leaf.frozen)
         return leaf
 
-    return jtu.tree_map(_replace_value, filter_spec, is_leaf=is_parameter)
+    return jax.tree.map(_replace_value, filter_spec, is_leaf=is_parameter)
 
 
 def partition(tree: PyTree) -> tuple[PyTree, PyTree]:
@@ -175,11 +174,11 @@ def sample(tree: PyTree, key: PRNGKeyArray) -> PyTree:
     """
     # partition the tree into parameters and the rest
     params_tree, rest_tree = eqx.partition(tree, is_parameter, is_leaf=is_parameter)
-    params_structure = jax.tree_util.tree_structure(params_tree)
+    params_structure = jax.tree.structure(params_tree)
     n_params = params_structure.num_leaves  # type: ignore[attr-defined]
 
     keys = jax.random.split(key, n_params)
-    keys_tree = jax.tree_util.tree_unflatten(params_structure, keys)
+    keys_tree = jax.tree.unflatten(params_structure, keys)
 
     def _sample(param: Parameter, key: Parameter) -> Array:
         if isinstance(param.prior, PDFLike):
@@ -209,7 +208,7 @@ def _sample(param: Parameter, key: Parameter) -> Array:
         return eqx.tree_at(lambda p: p.value, param, sampled_value)
 
     # sample for each parameter
-    sampled_params_tree = jtu.tree_map(
+    sampled_params_tree = jax.tree.map(
         _sample, params_tree, keys_tree, is_leaf=is_parameter
     )
 
@@ -246,7 +245,6 @@ def model(*parameters: PyTree[evm.Parameter]):
 
         # More general case of correlating any PyTree of parameters
         from typing import NamedTuple
-        import jax.tree_util as jtu
 
 
         class Params(NamedTuple):
@@ -256,11 +254,11 @@ class Params(NamedTuple):
         params = Params(mu=evm.Parameter(1.0), syst=evm.NormalParameter(0.0))
 
         def model(params: Params):
-            flat_params, tree_def = jtu.tree_flatten(params, evm.parameter.is_parameter)
+            flat_params, tree_def = jax.tree.flatten(params, evm.parameter.is_parameter)
 
             # correlate the parameters
             correlated_flat_params = evm.parameter.correlate(*flat_params)
-            correlated_params = jtu.tree_unflatten(tree_def, correlated_flat_params)
+            correlated_params = jax.tree.unflatten(tree_def, correlated_flat_params)
 
             # now correlated_params.mu and correlated_params.syst are correlated, i.e., they share the same value
             assert correlated_params.mu.value == correlated_params.syst.value

src/evermore/staterror.py

@@ -4,8 +4,8 @@
 from typing import cast
 
 import equinox as eqx
+import jax
 import jax.numpy as jnp
-import jax.tree_util as jtu
 from jaxtyping import Array, PyTree
 
 from evermore.custom_types import ModifierLike
@@ -67,7 +67,7 @@ def __init__(
         threshold: float = 10.0,
     ) -> None:
         assert (
-            jtu.tree_structure(hists) == jtu.tree_structure(histsw2)  # type: ignore[operator]
+            jax.tree.structure(hists) == jax.tree.structure(histsw2)  # type: ignore[operator]
         ), "The PyTree structure of hists and histsw2 must be the same!"
         self.hists = hists
         self.histsw2 = histsw2
@@ -80,10 +80,10 @@ def __init__(
 
         # setup params
         self.gaussians_global = NormalParameter(value=jnp.zeros_like(self.ntot))
-        self.gaussians_per_process = jtu.tree_map(
+        self.gaussians_per_process = jax.tree.map(
             lambda hist: NormalParameter(value=jnp.zeros_like(hist)), self.hists
         )
-        self.poissons_per_process = jtu.tree_map(
+        self.poissons_per_process = jax.tree.map(
             lambda w, w2: Parameter(
                 value=jnp.zeros_like(w),
                 prior=Poisson(