Merge pull request #14 from ucl-bug/dependabot/pip/jax-0.2.25

Bump jax from 0.2.21 to 0.2.25

Author: astanziola

_setup/requirements.txt

@@ -1,3 +1,3 @@
-jax==0.2.21
+jax==0.2.25
 jaxlib==0.1.71
 hashids==1.3.1

docs/notebooks/customizing_discretizations.ipynb

@@ -525,7 +525,7 @@
    ],
    "source": [
     "from jax import numpy as jnp\n",
-    "from jax.experimental import stax\n",
+    "from jax.example_libraries import stax\n",
     "\n",
     "init_random_params, predict = stax.serial(\n",
     "    stax.Dense(1024), stax.Relu,\n",
@@ -573,7 +573,7 @@
    ],
    "source": [
     "# Reusing the operator with a neural network discretization\n",
-    "from jax.experimental import stax\n",
+    "from jax.example_libraries import stax\n",
     "from functools import partial\n",
     "from jaxdf.discretization import Arbitrary\n",
     "\n",

docs/notebooks/example_1_paper.ipynb

@@ -121,7 +121,7 @@
     "from jaxdf.utils import join_dicts\n",
     "from jax import numpy as jnp\n",
     "from jax.scipy.sparse.linalg import gmres\n",
-    "from jax.experimental import optimizers\n",
+    "from jax.example_libraries import optimizers\n",
     "import jax\n",
     "\n",
     "# Settings\n",

docs/notebooks/helmholtz_pinn.ipynb

@@ -164,7 +164,7 @@
    "outputs": [],
    "source": [
     "from jaxdf.discretization import Arbitrary, UniformField\n",
-    "from jax.experimental import stax\n",
+    "from jax.example_libraries import stax\n",
     "from jax import random\n",
     "\n",
     "seed = random.PRNGKey(42)\n",
@@ -354,7 +354,7 @@
    ],
    "source": [
     "# Training loop\n",
-    "from jax.experimental import optimizers\n",
+    "from jax.example_libraries import optimizers\n",
     "from jax.tree_util import tree_multimap\n",
     "from jax import jit\n",
     "from tqdm import tqdm\n",

docs/notebooks/learn_fd_kernel.ipynb

@@ -283,7 +283,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from jax.experimental import optimizers\n",
+    "from jax.example_libraries import optimizers\n",
     "\n",
     "init_fun, update_fun, get_params = optimizers.adam(.001, b1=0.9, b2=0.9)\n",
     "opt_state = init_fun(jnp.asarray(k))"

docs/notebooks/pinn_burgers.ipynb

@@ -872,7 +872,7 @@
     }
    ],
    "source": [
-    "from jax.experimental import optimizers\n",
+    "from jax.example_libraries import optimizers\n",
     "from jax import jit\n",
     "from tqdm import tqdm\n",
     "\n",