From dc65594f8054e69583c0ffc668b8e6339ed8aebe Mon Sep 17 00:00:00 2001
From: Owen Lockwood <42878312+lockwo@users.noreply.github.com>
Date: Mon, 6 Jul 2026 16:35:06 -0700
Subject: [PATCH] fix

---
 distreqx/distributions/_distribution.py | 2 +-
 distreqx/distributions/_independent.py  | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/distreqx/distributions/_distribution.py b/distreqx/distributions/_distribution.py
index 2a8c80b..1d50dbd 100644
--- a/distreqx/distributions/_distribution.py
+++ b/distreqx/distributions/_distribution.py
@@ -66,7 +66,7 @@ def support(self) -> tuple[PyTree[Array], PyTree[Array]]:
     def dtype(self) -> jnp.dtype:
         """Data type of a sample"""
         sample_spec = jax.eval_shape(self.sample, jax.random.key(0))
-        return jax.tree_util.tree_map(lambda x: x.dtype, sample_spec)
+        return jax.tree.map(lambda x: x.dtype, sample_spec)
 
     @property
     def name(self) -> str:
diff --git a/distreqx/distributions/_independent.py b/distreqx/distributions/_independent.py
index 5f4dcb0..6964601 100644
--- a/distreqx/distributions/_independent.py
+++ b/distreqx/distributions/_independent.py
@@ -2,8 +2,8 @@
 
 import operator
 
+import jax
 import jax.numpy as jnp
-import jax.tree_util as jtu
 from jaxtyping import Array, Key, PyTree
 
 from ._distribution import (
@@ -15,8 +15,8 @@
 
 
 def _reduce_helper(pytree: PyTree) -> Array:
-    sum_over_leaves = jtu.tree_map(jnp.sum, pytree)
-    return jtu.tree_reduce(operator.add, sum_over_leaves)
+    sum_over_leaves = jax.tree.map(jnp.sum, pytree)
+    return jax.tree.reduce(operator.add, sum_over_leaves)
 
 
 class Independent(