jrandom->jr

Author: patrick-kidger

benchmarks/small_neural_ode.py

@@ -10,7 +10,7 @@
 import jax.experimental.ode as experimental
 import jax.nn as jnn
 import jax.numpy as jnp
-import jax.random as jrandom
+import jax.random as jr
 import numpy as np
 import torch  # pyright: ignore
 import torchdiffeq  # pyright: ignore
@@ -44,7 +44,7 @@ def __init__(self):
             depth=1,
             activation=jnn.softplus,
             final_activation=jnn.tanh,
-            key=jrandom.PRNGKey(0),
+            key=jr.PRNGKey(0),
         )
 
     def __call__(self, t, y, args):
@@ -182,7 +182,7 @@ def run(multiple, grad, batch_size=64, t1=100):
         func_torch[2].weight.copy_(torch.tensor(np.asarray(func_jax.layers[1].weight)))  # pyright: ignore
         func_torch[2].bias.copy_(torch.tensor(np.asarray(func_jax.layers[1].bias)))  # pyright: ignore
 
-    y0_jax = jrandom.normal(jrandom.PRNGKey(1), (batch_size, 4))
+    y0_jax = jr.normal(jr.PRNGKey(1), (batch_size, 4))
     y0_torch = torch.tensor(np.asarray(y0_jax))
 
     time_torch(neural_ode_torch, y0_torch, t1, grad)

diffrax/_brownian/path.py

@@ -4,7 +4,7 @@
 import equinox.internal as eqxi
 import jax
 import jax.numpy as jnp
-import jax.random as jrandom
+import jax.random as jr
 import jax.tree_util as jtu
 from jaxtyping import Array, PRNGKeyArray, PyTree
 
@@ -81,8 +81,8 @@ def evaluate(
         t1 = cast(RealScalarLike, t1)
         t0_ = force_bitcast_convert_type(t0, jnp.int32)
         t1_ = force_bitcast_convert_type(t1, jnp.int32)
-        key = jrandom.fold_in(self.key, t0_)
-        key = jrandom.fold_in(key, t1_)
+        key = jr.fold_in(self.key, t0_)
+        key = jr.fold_in(key, t1_)
         key = split_by_tree(key, self.shape)
         return jtu.tree_map(
             lambda key, shape: self._evaluate_leaf(t0, t1, key, shape), key, self.shape
@@ -91,7 +91,7 @@ def evaluate(
     def _evaluate_leaf(
         self, t0: RealScalarLike, t1: RealScalarLike, key, shape: jax.ShapeDtypeStruct
     ):
-        return jrandom.normal(key, shape.shape, shape.dtype) * jnp.sqrt(t1 - t0).astype(
+        return jr.normal(key, shape.shape, shape.dtype) * jnp.sqrt(t1 - t0).astype(
             shape.dtype
         )
 

diffrax/_brownian/tree.py

@@ -5,7 +5,7 @@
 import jax
 import jax.lax as lax
 import jax.numpy as jnp
-import jax.random as jrandom
+import jax.random as jr
 import jax.tree_util as jtu
 from jaxtyping import Array, Float, PRNGKeyArray, PyTree
 
@@ -112,7 +112,7 @@ def _brownian_bridge(self, s, t, u, w_s, w_u, key, shape, dtype):
         mean = w_s + (w_u - w_s) * ((t - s) / (u - s))
         var = (u - t) * (t - s) / (u - s)
         std = jnp.sqrt(var)
-        return mean + std * jrandom.normal(key, shape, dtype)
+        return mean + std * jr.normal(key, shape, dtype)
 
     def _evaluate_leaf(
         self,
@@ -140,9 +140,9 @@ def _evaluate_leaf(
         # errors are only raised after everything has finished executing.
         τ = jnp.clip(τ, t0, t1).astype(dtype)
 
-        key, init_key = jrandom.split(key, 2)
+        key, init_key = jr.split(key, 2)
         thalf = t0 + 0.5 * (t1 - t0)
-        w_t1 = jrandom.normal(init_key, shape, dtype) * jnp.sqrt(t1 - t0)
+        w_t1 = jr.normal(init_key, shape, dtype) * jnp.sqrt(t1 - t0)
         w_thalf = self._brownian_bridge(t0, thalf, t1, 0, w_t1, key, shape, dtype)
         init_state = _State(
             s=t0,
@@ -164,7 +164,7 @@ def _cond_fun(_state):
             return (_state.u - _state.s) > self.tol
 
         def _body_fun(_state):
-            _key1, _key2 = jrandom.split(_state.key, 2)
+            _key1, _key2 = jr.split(_state.key, 2)
             _cond = τ > _state.t
             _s = jnp.where(_cond, _state.t, _state.s)
             _u = jnp.where(_cond, _state.u, _state.t)

docs/usage/getting-started.md

@@ -69,13 +69,13 @@ $y(0) = 1 \qquad \mathrm{d}y(t) = -y(t)\mathrm{d}t + \frac{t}{10}\mathrm{d}w(t)$
 over the interval $[0, 3]$.
 
 ```python
-import jax.random as jrandom
+import jax.random as jr
 from diffrax import diffeqsolve, ControlTerm, Euler, MultiTerm, ODETerm, SaveAt, VirtualBrownianTree
 
 t0, t1 = 1, 3
 drift = lambda t, y, args: -y
 diffusion = lambda t, y, args: 0.1 * t
-brownian_motion = VirtualBrownianTree(t0, t1, tol=1e-3, shape=(), key=jrandom.PRNGKey(0))
+brownian_motion = VirtualBrownianTree(t0, t1, tol=1e-3, shape=(), key=jr.PRNGKey(0))
 terms = MultiTerm(ODETerm(drift), ControlTerm(diffusion, brownian_motion))
 solver = Euler()
 saveat = SaveAt(dense=True)

examples/continuous_normalising_flow.ipynb

@@ -73,7 +73,7 @@
     "import jax.lax as lax\n",
     "import jax.nn as jnn\n",
     "import jax.numpy as jnp\n",
-    "import jax.random as jrandom\n",
+    "import jax.random as jr\n",
     "import matplotlib.pyplot as plt\n",
     "import optax  # https://github.com/deepmind/optax\n",
     "import scipy.stats as stats\n",
@@ -111,7 +111,7 @@
     "\n",
     "    def __init__(self, *, data_size, width_size, depth, key, **kwargs):\n",
     "        super().__init__(**kwargs)\n",
-    "        keys = jrandom.split(key, depth + 1)\n",
+    "        keys = jr.split(key, depth + 1)\n",
     "        layers = []\n",
     "        if depth == 0:\n",
     "            layers.append(\n",
@@ -150,7 +150,7 @@
     "\n",
     "    def __init__(self, *, in_size, out_size, key, **kwargs):\n",
     "        super().__init__(**kwargs)\n",
-    "        key1, key2, key3 = jrandom.split(key, 3)\n",
+    "        key1, key2, key3 = jr.split(key, 3)\n",
     "        self.lin1 = eqx.nn.Linear(in_size, out_size, key=key1)\n",
     "        self.lin2 = eqx.nn.Linear(1, out_size, key=key2)\n",
     "        self.lin3 = eqx.nn.Linear(1, out_size, use_bias=False, key=key3)\n",
@@ -251,7 +251,7 @@
     "        **kwargs,\n",
     "    ):\n",
     "        super().__init__(**kwargs)\n",
-    "        keys = jrandom.split(key, num_blocks)\n",
+    "        keys = jr.split(key, num_blocks)\n",
     "        self.funcs = [\n",
     "            Func(\n",
     "                data_size=data_size,\n",
@@ -274,7 +274,7 @@
     "        else:\n",
     "            term = diffrax.ODETerm(approx_logp_wrapper)\n",
     "        solver = diffrax.Tsit5()\n",
-    "        eps = jrandom.normal(key, y.shape)\n",
+    "        eps = jr.normal(key, y.shape)\n",
     "        delta_log_likelihood = 0.0\n",
     "        for func in reversed(self.funcs):\n",
     "            y = (y, delta_log_likelihood)\n",
@@ -286,7 +286,7 @@
     "\n",
     "    # Runs forward-in-time to draw samples from the CNF.\n",
     "    def sample(self, *, key):\n",
-    "        y = jrandom.normal(key, (self.data_size,))\n",
+    "        y = jr.normal(key, (self.data_size,))\n",
     "        for func in self.funcs:\n",
     "            term = diffrax.ODETerm(func)\n",
     "            solver = diffrax.Tsit5()\n",
@@ -300,7 +300,7 @@
     "        t_so_far = self.t0\n",
     "        t_end = self.t0 + (self.t1 - self.t0) * len(self.funcs)\n",
     "        save_times = jnp.linspace(self.t0, t_end, 6)\n",
-    "        y = jrandom.normal(key, (self.data_size,))\n",
+    "        y = jr.normal(key, (self.data_size,))\n",
     "        out = []\n",
     "        for i, func in enumerate(self.funcs):\n",
     "            if i == len(self.funcs) - 1:\n",
@@ -404,7 +404,7 @@
     "class DataLoader(eqx.Module):\n",
     "    arrays: tuple[jnp.ndarray, ...]\n",
     "    batch_size: int\n",
-    "    key: jrandom.PRNGKey\n",
+    "    key: jr.PRNGKey\n",
     "\n",
     "    def __check_init__(self):\n",
     "        dataset_size = self.arrays[0].shape[0]\n",
@@ -414,8 +414,8 @@
     "        dataset_size = self.arrays[0].shape[0]\n",
     "        num_batches = dataset_size // self.batch_size\n",
     "        epoch = step // num_batches\n",
-    "        key = jrandom.fold_in(self.key, epoch)\n",
-    "        perm = jrandom.permutation(key, jnp.arange(dataset_size))\n",
+    "        key = jr.fold_in(self.key, epoch)\n",
+    "        perm = jr.permutation(key, jnp.arange(dataset_size))\n",
     "        start = (step % num_batches) * self.batch_size\n",
     "        slice_size = self.batch_size\n",
     "        batch_indices = lax.dynamic_slice_in_dim(perm, start, slice_size)\n",
@@ -464,8 +464,8 @@
     "    else:\n",
     "        out_path = pathlib.Path(out_path)\n",
     "\n",
-    "    key = jrandom.PRNGKey(seed)\n",
-    "    model_key, loader_key, loss_key, sample_key = jrandom.split(key, 4)\n",
+    "    key = jr.PRNGKey(seed)\n",
+    "    model_key, loader_key, loss_key, sample_key = jr.split(key, 4)\n",
     "\n",
     "    dataset, weights, mean, std, img, width, height = get_data(in_path)\n",
     "    dataset_size, data_size = dataset.shape\n",
@@ -486,9 +486,9 @@
     "    @eqx.filter_value_and_grad\n",
     "    def loss(model, data, weight, loss_key):\n",
     "        batch_size, _ = data.shape\n",
-    "        noise_key, train_key = jrandom.split(loss_key, 2)\n",
-    "        train_key = jrandom.split(key, batch_size)\n",
-    "        data = data + jrandom.normal(noise_key, data.shape) * 0.5 / std\n",
+    "        noise_key, train_key = jr.split(loss_key, 2)\n",
+    "        train_key = jr.split(key, batch_size)\n",
+    "        data = data + jr.normal(noise_key, data.shape) * 0.5 / std\n",
     "        log_likelihood = jax.vmap(model.train)(data, key=train_key)\n",
     "        return -jnp.mean(weight * log_likelihood)  # minimise negative log-likelihood\n",
     "\n",
@@ -514,7 +514,7 @@
     "            value = value + value_\n",
     "            grads = jax.tree_util.tree_map(lambda a, b: a + b, grads, grads_)\n",
     "            step = step + 1\n",
-    "            loss_key = jrandom.split(loss_key, 1)[0]\n",
+    "            loss_key = jr.split(loss_key, 1)[0]\n",
     "            return value, grads, step, loss_key\n",
     "\n",
     "        value, grads, step, loss_key = lax.fori_loop(\n",
@@ -537,7 +537,7 @@
     "            print(f\"Step: {step}, Loss: {value}, Computation time: {end - start}\")\n",
     "\n",
     "    num_samples = 5000\n",
-    "    sample_key = jrandom.split(sample_key, num_samples)\n",
+    "    sample_key = jr.split(sample_key, num_samples)\n",
     "    samples = jax.vmap(model.sample)(key=sample_key)\n",
     "    sample_flows = jax.vmap(model.sample_flow, out_axes=-1)(key=sample_key)\n",
     "    fig, (*axs, ax, axtrue) = plt.subplots(\n",

examples/latent_ode.ipynb

@@ -78,7 +78,7 @@
     "import jax\n",
     "import jax.nn as jnn\n",
     "import jax.numpy as jnp\n",
-    "import jax.random as jrandom\n",
+    "import jax.random as jr\n",
     "import matplotlib\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
@@ -142,7 +142,7 @@
     "    ):\n",
     "        super().__init__(**kwargs)\n",
     "\n",
-    "        mkey, gkey, hlkey, lhkey, hdkey = jrandom.split(key, 5)\n",
+    "        mkey, gkey, hlkey, lhkey, hdkey = jr.split(key, 5)\n",
     "\n",
     "        scale = jnp.ones(())\n",
     "        mlp = eqx.nn.MLP(\n",
@@ -175,7 +175,7 @@
     "        context = self.hidden_to_latent(hidden)\n",
     "        mean, logstd = context[: self.latent_size], context[self.latent_size :]\n",
     "        std = jnp.exp(logstd)\n",
-    "        latent = mean + jrandom.normal(key, (self.latent_size,)) * std\n",
+    "        latent = mean + jr.normal(key, (self.latent_size,)) * std\n",
     "        return latent, mean, std\n",
     "\n",
     "    # Decoder of the VAE\n",
@@ -209,7 +209,7 @@
     "\n",
     "    # Run just the decoder during inference.\n",
     "    def sample(self, ts, *, key):\n",
-    "        latent = jrandom.normal(key, (self.latent_size,))\n",
+    "        latent = jr.normal(key, (self.latent_size,))\n",
     "        return self._sample(ts, latent)"
    ]
   },
@@ -231,13 +231,13 @@
    "outputs": [],
    "source": [
     "def get_data(dataset_size, *, key):\n",
-    "    ykey, tkey1, tkey2 = jrandom.split(key, 3)\n",
+    "    ykey, tkey1, tkey2 = jr.split(key, 3)\n",
     "\n",
-    "    y0 = jrandom.normal(ykey, (dataset_size, 2))\n",
+    "    y0 = jr.normal(ykey, (dataset_size, 2))\n",
     "\n",
     "    t0 = 0\n",
-    "    t1 = 2 + jrandom.uniform(tkey1, (dataset_size,))\n",
-    "    ts = jrandom.uniform(tkey2, (dataset_size, 20)) * (t1[:, None] - t0) + t0\n",
+    "    t1 = 2 + jr.uniform(tkey1, (dataset_size,))\n",
+    "    ts = jr.uniform(tkey2, (dataset_size, 20)) * (t1[:, None] - t0) + t0\n",
     "    ts = jnp.sort(ts)\n",
     "    dt0 = 0.1\n",
     "\n",
@@ -273,8 +273,8 @@
     "    assert all(array.shape[0] == dataset_size for array in arrays)\n",
     "    indices = jnp.arange(dataset_size)\n",
     "    while True:\n",
-    "        perm = jrandom.permutation(key, indices)\n",
-    "        (key,) = jrandom.split(key, 1)\n",
+    "        perm = jr.permutation(key, indices)\n",
+    "        (key,) = jr.split(key, 1)\n",
     "        start = 0\n",
     "        end = batch_size\n",
     "        while start < dataset_size:\n",
@@ -311,8 +311,8 @@
     "    depth=2,\n",
     "    seed=5678,\n",
     "):\n",
-    "    key = jrandom.PRNGKey(seed)\n",
-    "    data_key, model_key, loader_key, train_key, sample_key = jrandom.split(key, 5)\n",
+    "    key = jr.PRNGKey(seed)\n",
+    "    data_key, model_key, loader_key, train_key, sample_key = jr.split(key, 5)\n",
     "\n",
     "    ts, ys = get_data(dataset_size, key=data_key)\n",
     "\n",
@@ -328,14 +328,14 @@
     "    @eqx.filter_value_and_grad\n",
     "    def loss(model, ts_i, ys_i, key_i):\n",
     "        batch_size, _ = ts_i.shape\n",
-    "        key_i = jrandom.split(key_i, batch_size)\n",
+    "        key_i = jr.split(key_i, batch_size)\n",
     "        loss = jax.vmap(model.train)(ts_i, ys_i, key=key_i)\n",
     "        return jnp.mean(loss)\n",
     "\n",
     "    @eqx.filter_jit\n",
     "    def make_step(model, opt_state, ts_i, ys_i, key_i):\n",
     "        value, grads = loss(model, ts_i, ys_i, key_i)\n",
-    "        key_i = jrandom.split(key_i, 1)[0]\n",
+    "        key_i = jr.split(key_i, 1)[0]\n",
     "        updates, opt_state = optim.update(grads, opt_state)\n",
     "        model = eqx.apply_updates(model, updates)\n",
     "        return value, model, opt_state, key_i\n",

examples/neural_cde.ipynb

@@ -61,7 +61,7 @@
     "import jax\n",
     "import jax.nn as jnn\n",
     "import jax.numpy as jnp\n",
-    "import jax.random as jrandom\n",
+    "import jax.random as jr\n",
     "import jax.scipy as jsp\n",
     "import matplotlib\n",
     "import matplotlib.pyplot as plt\n",
@@ -136,7 +136,7 @@
     "\n",
     "    def __init__(self, data_size, hidden_size, width_size, depth, *, key, **kwargs):\n",
     "        super().__init__(**kwargs)\n",
-    "        ikey, fkey, lkey = jrandom.split(key, 3)\n",
+    "        ikey, fkey, lkey = jr.split(key, 3)\n",
     "        self.initial = eqx.nn.MLP(data_size, hidden_size, width_size, depth, key=ikey)\n",
     "        self.func = Func(data_size, hidden_size, width_size, depth, key=fkey)\n",
     "        self.linear = eqx.nn.Linear(hidden_size, 1, key=lkey)\n",
@@ -195,9 +195,9 @@
    "outputs": [],
    "source": [
     "def get_data(dataset_size, add_noise, *, key):\n",
-    "    theta_key, noise_key = jrandom.split(key, 2)\n",
+    "    theta_key, noise_key = jr.split(key, 2)\n",
     "    length = 100\n",
-    "    theta = jrandom.uniform(theta_key, (dataset_size,), minval=0, maxval=2 * math.pi)\n",
+    "    theta = jr.uniform(theta_key, (dataset_size,), minval=0, maxval=2 * math.pi)\n",
     "    y0 = jnp.stack([jnp.cos(theta), jnp.sin(theta)], axis=-1)\n",
     "    ts = jnp.broadcast_to(jnp.linspace(0, 4 * math.pi, length), (dataset_size, length))\n",
     "    matrix = jnp.array([[-0.3, 2], [-2, -0.3]])\n",
@@ -207,7 +207,7 @@
     "    ys = jnp.concatenate([ts[:, :, None], ys], axis=-1)  # time is a channel\n",
     "    ys = ys.at[: dataset_size // 2, :, 1].multiply(-1)\n",
     "    if add_noise:\n",
-    "        ys = ys + jrandom.normal(noise_key, ys.shape) * 0.1\n",
+    "        ys = ys + jr.normal(noise_key, ys.shape) * 0.1\n",
     "    coeffs = jax.vmap(diffrax.backward_hermite_coefficients)(ts, ys)\n",
     "    labels = jnp.zeros((dataset_size,))\n",
     "    labels = labels.at[: dataset_size // 2].set(1.0)\n",
@@ -227,8 +227,8 @@
     "    assert all(array.shape[0] == dataset_size for array in arrays)\n",
     "    indices = jnp.arange(dataset_size)\n",
     "    while True:\n",
-    "        perm = jrandom.permutation(key, indices)\n",
-    "        (key,) = jrandom.split(key, 1)\n",
+    "        perm = jr.permutation(key, indices)\n",
+    "        (key,) = jr.split(key, 1)\n",
     "        start = 0\n",
     "        end = batch_size\n",
     "        while end < dataset_size:\n",
@@ -264,8 +264,8 @@
     "    depth=1,\n",
     "    seed=5678,\n",
     "):\n",
-    "    key = jrandom.PRNGKey(seed)\n",
-    "    train_data_key, test_data_key, model_key, loader_key = jrandom.split(key, 4)\n",
+    "    key = jr.PRNGKey(seed)\n",
+    "    train_data_key, test_data_key, model_key, loader_key = jr.split(key, 4)\n",
     "\n",
     "    ts, coeffs, labels, data_size = get_data(\n",
     "        dataset_size, add_noise, key=train_data_key\n",