Fix import path for tree_map

`jax.tree_map` has been deprecated, see: https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-26-april-3-2024
probml · Jul 15, 2024 · 28de6b0 · 28de6b0
1 parent fceca0e
commit 28de6b0
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/dynamax/hidden_markov_model/models/arhmm.py b/dynamax/hidden_markov_model/models/arhmm.py
@@ -1,6 +1,7 @@
 import jax.numpy as jnp
 import jax.random as jr
-from jax import lax, tree_map
+from jax import lax
+from jax.tree_util import tree_map
 from jaxtyping import Float, Array
 from dynamax.hidden_markov_model.models.abstractions import HMM, HMMParameterSet, HMMPropertySet
 from dynamax.hidden_markov_model.models.initial import StandardHMMInitialState, ParamsStandardHMMInitialState

diff --git a/dynamax/utils/distributions_test.py b/dynamax/utils/distributions_test.py
@@ -1,7 +1,7 @@
 import pytest
 import jax.numpy as jnp
 import jax.random as jr
-from jax import tree_map
+from jax.tree_util import tree_map
 from jax.scipy.stats import norm
 from scipy.stats import invgamma
 from tensorflow_probability.substrates import jax as tfp

diff --git a/dynamax/utils/optimize.py b/dynamax/utils/optimize.py
@@ -1,7 +1,8 @@
 import jax.numpy as jnp
 import jax.random as jr
 import optax
-from jax import lax, value_and_grad, tree_map
+from jax import lax, value_and_grad
+from jax.tree_util import tree_map
 from dynamax.utils.utils import pytree_len