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

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",

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

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

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