deps: unpin jax/jaxlib now that numpyro 0.21.0 ships fix

examples/advanced/pymdp_with_neural_encoder.ipynb

@@ -274,7 +274,7 @@
    "source": [
     "def sample_categorical_batch(key: jnp.ndarray, probs: jnp.ndarray) -> jnp.ndarray:\n",
     "    keys = jr.split(key, probs.shape[0])\n",
-    "    return jax.vmap(lambda k, p: jr.categorical(k, jnp.log(jnp.clip(p, a_min=EPS))))(keys, probs).astype(jnp.int32)\n",
+    "    return jax.vmap(lambda k, p: jr.categorical(k, jnp.log(jnp.clip(p, min=EPS))))(keys, probs).astype(jnp.int32)\n",
     "\n",
     "\n",
     "def sample_balanced_initial_states(key: jnp.ndarray, num_sequences: int) -> jnp.ndarray:\n",
@@ -402,7 +402,7 @@
     "def model_obs_categorical(model: FrontendModel, x_t: jnp.ndarray):\n",
     "    logits = jax.vmap(model.encoder)(x_t) / LOGIT_TEMPERATURE\n",
     "    probs = jnn.softmax(logits, axis=-1)\n",
-    "    probs = jnp.clip(probs, a_min=EPS)\n",
+    "    probs = jnp.clip(probs, min=EPS)\n",
     "    probs = probs / probs.sum(axis=-1, keepdims=True)\n",
     "    return [probs]\n",
     "\n",
@@ -444,10 +444,10 @@
     "        return jnp.array(0.0)\n",
     "\n",
     "    pred_obs_next = jnp.einsum('btk,ok->bto', pred_next_state_seq, A_FIXED_SINGLE)\n",
-    "    pred_obs_next = jnp.clip(pred_obs_next, a_min=EPS)\n",
+    "    pred_obs_next = jnp.clip(pred_obs_next, min=EPS)\n",
     "    pred_obs_next = pred_obs_next / pred_obs_next.sum(axis=-1, keepdims=True)\n",
     "\n",
-    "    obs_next = jnp.clip(obs_seq[:, 1:, :], a_min=EPS)\n",
+    "    obs_next = jnp.clip(obs_seq[:, 1:, :], min=EPS)\n",
     "    kl = jnp.sum(obs_next * (jnp.log(obs_next) - jnp.log(pred_obs_next)), axis=-1)\n",
     "    return jnp.mean(kl)\n",
     "\n",
@@ -710,7 +710,7 @@
     "    t = true_states.reshape(-1)\n",
     "    p = pred_states.reshape(-1)\n",
     "    cm = cm.at[t, p].add(1)\n",
-    "    cm = cm / jnp.clip(cm.sum(axis=1, keepdims=True), a_min=1.0)\n",
+    "    cm = cm / jnp.clip(cm.sum(axis=1, keepdims=True), min=1.0)\n",
     "    return cm\n",
     "\n",
     "\n",

pyproject.toml

@@ -24,11 +24,8 @@ classifiers=[
 requires-python = ">=3.10"
 dependencies = [
     'numpy>=1.19.5',
-    # jax/jaxlib upper-bounded until numpyro ships the fix from
-    # https://github.com/pyro-ppl/numpyro/pull/2173 (jax 0.10.0 removed
-    # xla_pmap_p, which breaks numpyro -> pybefit imports). See issue #389.
-    'jax>=0.3.4,<0.10',
-    'jaxlib>=0.3.4,<0.10',
+    'jax>=0.3.4',
+    'jaxlib>=0.3.4',
     'equinox>=0.9',
     'multimethod>=1.11',
     'matplotlib>=3.1.3',
@@ -60,7 +57,7 @@ test = [
 
 [project.optional-dependencies]
 gpu = [
-    'jax[cuda12]>=0.3.4,<0.10',
+    'jax[cuda12]>=0.3.4',
 ]
 docs = [
     "mkdocs>=1.6.1,<2",

setup.cfg

@@ -24,12 +24,8 @@ project_urls =
 packages = find:
 install_requires =
     numpy>=1.19.5
-    # jax/jaxlib upper-bounded until numpyro ships the fix from
-    # https://github.com/pyro-ppl/numpyro/pull/2173 (jax 0.10.0 removed
-    # xla_pmap_p, which breaks numpyro -> pybefit imports). See issue #389.
-    # Keep in sync with pyproject.toml.
-    jax>=0.3.4,<0.10
-    jaxlib>=0.3.4,<0.10
+    jax>=0.3.4
+    jaxlib>=0.3.4
     equinox>=0.9
     multimethod>=1.11
     matplotlib>=3.1.3
@@ -40,7 +36,7 @@ include_package_data = True
 
 [options.extras_require]
 gpu =
-    jax[cuda12]>=0.3.4,<0.10
+    jax[cuda12]>=0.3.4
 
 [options.package_data]
 pymdp = envs/assets/*