Merge pull request #369 from probml/ghd/fix-tree-map-import

Fix import path for `tree_map`
probml · Jul 15, 2024 · be7212c · be7212c
2 parents fceca0e + 0615c95
commit be7212c
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 4 deletions.
diff --git a/docs/notebooks/slds/rbpf_maneuver.ipynb b/docs/notebooks/slds/rbpf_maneuver.ipynb
@@ -32,7 +32,8 @@
     "from functools import partial\n",
     "import sys \n",
     "sys.path.append('/Users/kostastsampourakis/Desktop/code/Python/projects/dynamax')\n",
-    "from jax import vmap, tree_map, jit\n",
+    "from jax import vmap, jit\n",
+    "from jax.tree_util import tree_map\n",
     "from dynamax.slds.inference import ParamsSLDS, LGParamsSLDS, DiscreteParamsSLDS, rbpfilter, rbpfilter_optimal\n",
     "from dynamax.slds.models import SLDS\n",
     "# import MVN from tfd\n",

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