From efee0fbf7abe3116cba10252e7d5cbb06a095b1b Mon Sep 17 00:00:00 2001 From: Jan-Williams Date: Wed, 20 May 2026 09:13:31 -0700 Subject: [PATCH 1/2] Updated jit compatibility example --- README.md | 3 + examples/jit_compatibility.ipynb | 519 +++++++++++++++++++++++++++++++ mkdocs.yml | 1 + src/orc/classifier/base.py | 14 +- src/orc/classifier/models.py | 5 +- src/orc/control/base.py | 12 +- src/orc/control/models.py | 7 +- src/orc/drivers.py | 24 +- src/orc/embeddings.py | 30 +- src/orc/forecaster/base.py | 12 +- src/orc/forecaster/models.py | 13 +- src/orc/readouts.py | 40 +-- 12 files changed, 603 insertions(+), 77 deletions(-) create mode 100644 examples/jit_compatibility.ipynb diff --git a/README.md b/README.md index 2201c37..507079d 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,9 @@ plt.show() ORC Logo +## jit, vmap, grad... + +ORC models are built on top of [Equinox](https://docs.kidger.site/equinox/), and as a result we strongly recommend the use of Equinox transforms `eqx.filter_{jit, vmap, grad}` over `jax.{jit, vmap, grad}`. For more details, please check out the [JAX JIT Compatibility](https://Jan-Williams.github.io/OpenReservoirComputing/examples/jit_compatibility/) example notebook. ## Contribution guidelines First off, thanks for helping out! We appreciate your willingness to contribute! To get started, clone the repo and install the developer dependencies of ORC. diff --git a/examples/jit_compatibility.ipynb b/examples/jit_compatibility.ipynb new file mode 100644 index 0000000..8d575f9 --- /dev/null +++ b/examples/jit_compatibility.ipynb @@ -0,0 +1,519 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a1b2c3d4", + "metadata": {}, + "source": [ + "# JAX JIT Compatibility\n", + "\n", + "## TL;DR\n", + "Use `eqx.filter_{vmap, grad, jit}` over `jax.{vmap, grad, jit}` unless you *__know__* otherwise.\n", + "\n", + "\n", + "ORC models are [Equinox](https://docs.kidger.site/equinox/) modules, which means they are registered as JAX pytrees. This notebook covers the sharp bits around using `jax.jit` and `eqx.filter_jit` with ORC models and shows the correct pattern for each case.\n", + "\n", + "## Why eqx.filter_{jit, vmap, grad}\n", + "\n", + "Many ORC layers use Jax sparse `BCOO` matrices to accelerate computations. `BCOO` matrices are **unhashable**. When you write `jax.jit(esn.forecast)`, JAX closes over `esn` and tries to hash it as a static compile-time constant, which fails because it contains a `BCOO`. \n", + "\n", + "`eqx.filter_jit` avoids this entirely by treating the model as a pytree. Array leaves (including `BCOO`) are **traced**, and non-array structural fields are treated as **static**. No hashing is required. Since all ORC inference and training methods are already decorated with `@eqx.filter_jit`, calling them directly compiles on the first call — you rarely need to add `eqx.filter_jit` explicitly.\n", + "\n", + "## When `jax.jit` works\n", + "\n", + "If you need fine-grained control over which arguments are static for recompilation purposes, `jax.jit` can work provided you:\n", + "\n", + "1. Pass the model as an **explicit pytree argument** (not closed over), and\n", + "2. Mark integer scalar arguments that appear as `jax.lax.scan` loop lengths as **static** via `static_argnums` / `static_argnames`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "b2c3d4e5", + "metadata": {}, + "outputs": [], + "source": [ + "import functools\n", + "\n", + "import equinox as eqx\n", + "import jax\n", + "import jax.numpy as jnp\n", + "\n", + "jax.config.update(\"jax_enable_x64\", True)\n", + "\n", + "import orc" + ] + }, + { + "cell_type": "markdown", + "id": "c3d4e5f6", + "metadata": {}, + "source": [ + "## 1. Forecasting\n", + "\n", + "Train an `ESNForecaster` on the Lorenz-63 system." + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "id": "d4e5f6a7", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Reservoir states shape: (2400, 1, 500)\n" + ] + } + ], + "source": [ + "U, t = orc.data.lorenz63(tN=30, dt=0.01)\n", + "split = int(0.8 * len(U))\n", + "U_train, U_test = U[:split], U[split:]\n", + "\n", + "esn = orc.forecaster.ESNForecaster(data_dim=3, res_dim=500, seed=0)\n", + "esn, R = orc.forecaster.train_RCForecaster(esn, U_train)\n", + "fcast_len = len(U_test)\n", + "print(f\"Reservoir states shape: {R.shape}\")" + ] + }, + { + "cell_type": "markdown", + "id": "e5f6a7b8", + "metadata": {}, + "source": [ + "### Pattern 1 — closed-over `jax.jit` (does **not** work)\n", + "\n", + "`jax.jit(esn.forecast)` closes over `esn`. JAX tries to hash the model as a compile-time static constant, which fails because `wr` is an unhashable `BCOO` sparse matrix." + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "id": "f6a7b8c9", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "TypeError: unhashable type: 'BCOO'\n" + ] + } + ], + "source": [ + "try:\n", + " U_pred = jax.jit(esn.forecast)(fcast_len=fcast_len, res_state=R[-1])\n", + "except TypeError as e:\n", + " print(f\"TypeError: {e}\")" + ] + }, + { + "cell_type": "markdown", + "id": "a7b8c9d0", + "metadata": {}, + "source": [ + "### Pattern 2 — `eqx.filter_jit` (recommended)\n", + "\n", + "`eqx.filter_jit` treats the model as a pytree, tracing its array leaves and keeping non-array fields static. No hashing is ever attempted.\n", + "\n", + "Because `forecast` is already decorated with `@eqx.filter_jit` internally, calling `esn.forecast(...)` directly compiles on first use — wrapping it again in `eqx.filter_jit` is equivalent and harmless." + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "id": "b8c9d0e1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Forecast shape: (600, 3)\n", + "Forecast shape (explicit filter_jit): (600, 3)\n" + ] + } + ], + "source": [ + "# Calling the method directly compiles on first use (already @eqx.filter_jit)\n", + "U_pred = esn.forecast(fcast_len=fcast_len, res_state=R[-1])\n", + "print(f\"Forecast shape: {U_pred.shape}\")\n", + "\n", + "# Explicit eqx.filter_jit wrap is equivalent:\n", + "U_pred = eqx.filter_jit(esn.forecast)(fcast_len=fcast_len, res_state=R[-1])\n", + "print(f\"Forecast shape (explicit filter_jit): {U_pred.shape}\")" + ] + }, + { + "cell_type": "markdown", + "id": "c9d0e1f2", + "metadata": {}, + "source": [ + "### Pattern 3 — `jax.jit` with explicit pytree argument\n", + "\n", + "Pass the model as an explicit argument so JAX traces through it as a pytree instead of hashing it. The `fcast_len` argument is used as a `jax.lax.scan` loop length, so it must be a concrete value at trace time — mark it static." + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "id": "d0e1f2a3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Forecast shape: (600, 3)\n" + ] + } + ], + "source": [ + "@functools.partial(jax.jit, static_argnums=(1,)) # fcast_len must be static\n", + "def jit_forecast(model, fcast_len, res_state):\n", + " return model.forecast(fcast_len, res_state)\n", + "\n", + "U_pred = jit_forecast(esn, fcast_len, R[-1])\n", + "print(f\"Forecast shape: {U_pred.shape}\")" + ] + }, + { + "cell_type": "markdown", + "id": "e1f2a3b4", + "metadata": {}, + "source": [ + "### Training with `jax.jit`\n", + "\n", + "`train_RCForecaster` and similar training functions work with `jax.jit` when the model is passed as an explicit pytree argument. Default scalar parameters (`spinup=0`, `batch_size=None`) remain as concrete Python values and do not need special handling. If you pass them explicitly, declare them static." + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "id": "f2a3b4c5", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Training with jax.jit (defaults): OK\n", + "Training with jax.jit (explicit spinup): OK\n" + ] + } + ], + "source": [ + "esn_fresh = orc.forecaster.ESNForecaster(data_dim=3, res_dim=500, seed=1)\n", + "\n", + "# Default arguments — works directly\n", + "esn_trained, R_trained = jax.jit(orc.forecaster.train_RCForecaster)(esn_fresh, U_train)\n", + "print(\"Training with jax.jit (defaults): OK\")\n", + "\n", + "# Non-default spinup — must be declared static (appears in a Python conditional)\n", + "train_fn = jax.jit(\n", + " orc.forecaster.train_RCForecaster,\n", + " static_argnames=(\"spinup\",),\n", + ")\n", + "esn_trained, R_trained = train_fn(esn_fresh, U_train, spinup=50)\n", + "print(\"Training with jax.jit (explicit spinup): OK\")" + ] + }, + { + "cell_type": "markdown", + "id": "a3b4c5d6", + "metadata": {}, + "source": [ + "## 2. Classification\n", + "\n", + "The same rules apply to `ESNClassifier`. Generate a simple multi-class sinusoidal dataset and train a classifier." + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "id": "b4c5d6e7", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Classifier trained.\n" + ] + } + ], + "source": [ + "n_classes, data_dim, seq_len = 3, 3, 200\n", + "key = jax.random.PRNGKey(42)\n", + "\n", + "def make_sequences(key, n_per_class):\n", + " t = jnp.linspace(0, 4 * jnp.pi, seq_len).reshape(-1, 1)\n", + " seqs, labels = [], []\n", + " for cls in range(n_classes):\n", + " freq = (cls + 1) * 0.5\n", + " for _ in range(n_per_class):\n", + " key, sub = jax.random.split(key)\n", + " noise = 0.05 * jax.random.normal(sub, (seq_len, data_dim))\n", + " seqs.append(jnp.sin(freq * t * jnp.arange(1, data_dim + 1)) + noise)\n", + " labels.append(cls)\n", + " return jnp.stack(seqs), jnp.array(labels, dtype=jnp.int32)\n", + "\n", + "train_seqs, train_labels = make_sequences(key, 20)\n", + "test_seqs, test_labels = make_sequences(jax.random.PRNGKey(99), 10)\n", + "\n", + "clf = orc.classifier.ESNClassifier(data_dim=data_dim, n_classes=n_classes, res_dim=500, seed=0)\n", + "clf = orc.classifier.train_RCClassifier(clf, train_seqs, train_labels)\n", + "print(\"Classifier trained.\")" + ] + }, + { + "cell_type": "markdown", + "id": "c5d6e7f8", + "metadata": {}, + "source": [ + "### Pattern 1 — closed-over `jax.jit` (does **not** work)" + ] + }, + { + "cell_type": "code", + "execution_count": 34, + "id": "d6e7f8a9", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "TypeError: unhashable type: 'BCOO'\n" + ] + } + ], + "source": [ + "try:\n", + " probs = jax.jit(clf.classify)(test_seqs[0])\n", + "except TypeError as e:\n", + " print(f\"TypeError: {e}\")" + ] + }, + { + "cell_type": "markdown", + "id": "e7f8a9b0", + "metadata": {}, + "source": [ + "### Pattern 2 — `eqx.filter_jit` (recommended)" + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "id": "f8a9b0c1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Predicted class: 0, True class: 0\n", + "Test accuracy: 100.0%\n" + ] + } + ], + "source": [ + "# Single sample\n", + "probs = clf.classify(test_seqs[0])\n", + "print(f\"Predicted class: {jnp.argmax(probs)}, True class: {test_labels[0]}\")\n", + "\n", + "# Batch over test set with vmap\n", + "all_probs = jax.vmap(clf.classify)(test_seqs)\n", + "accuracy = jnp.mean(jnp.argmax(all_probs, axis=1) == test_labels)\n", + "print(f\"Test accuracy: {accuracy:.1%}\")" + ] + }, + { + "cell_type": "markdown", + "id": "a9b0c1d2", + "metadata": {}, + "source": [ + "### Pattern 3 — `jax.jit` with explicit pytree argument\n", + "\n", + "`classify` has no integer loop-length arguments, so no `static_argnums` is needed — just pass the model explicitly." + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "id": "b0c1d2e3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Predicted class: 0, True class: 0\n" + ] + } + ], + "source": [ + "@jax.jit\n", + "def jit_classify(model, in_seq):\n", + " return model.classify(in_seq)\n", + "\n", + "probs = jit_classify(clf, test_seqs[0])\n", + "print(f\"Predicted class: {jnp.argmax(probs)}, True class: {test_labels[0]}\")" + ] + }, + { + "cell_type": "markdown", + "id": "c1d2e3f4", + "metadata": {}, + "source": [ + "## 3. Control\n", + "\n", + "The same rules apply to `ESNController`. We use a simple 2D linear dynamical system as a plant: two states, one scalar control input." + ] + }, + { + "cell_type": "code", + "execution_count": 37, + "id": "d2e3f4a5", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Controller trained. Reservoir states shape: (599, 500)\n" + ] + } + ], + "source": [ + "# Simple stable 2D linear plant\n", + "A_plant = jnp.array([[0.95, 0.1], [-0.1, 0.95]])\n", + "B_plant = jnp.array([[0.1], [0.0]])\n", + "\n", + "def simulate(u_seq, x0=None):\n", + " x = x0 if x0 is not None else jnp.zeros(2)\n", + " def step(x, u):\n", + " x_next = A_plant @ x + B_plant @ jnp.atleast_1d(u)\n", + " return x_next, x_next\n", + " _, states = jax.lax.scan(step, x, u_seq)\n", + " return states\n", + "\n", + "key = jax.random.PRNGKey(0)\n", + "u_train = jax.random.normal(key, (600, 1)) * 0.3\n", + "train_outputs = simulate(u_train)\n", + "\n", + "ctrl = orc.control.ESNController(data_dim=2, control_dim=1, res_dim=500, seed=0)\n", + "ctrl, R_ctrl = orc.control.train_RCController(ctrl, train_outputs[:-1], u_train[1:])\n", + "print(f\"Controller trained. Reservoir states shape: {R_ctrl.shape}\")" + ] + }, + { + "cell_type": "markdown", + "id": "e3f4a5b6", + "metadata": {}, + "source": [ + "### Pattern 1 — closed-over `jax.jit` (does **not** work)" + ] + }, + { + "cell_type": "code", + "execution_count": 38, + "id": "f4a5b6c7", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "TypeError: unhashable type: 'BCOO'\n" + ] + } + ], + "source": [ + "control_guess = jnp.zeros((100, 1))\n", + "\n", + "try:\n", + " output_traj = jax.jit(ctrl.apply_control)(control_guess, R_ctrl[-1])\n", + "except TypeError as e:\n", + " print(f\"TypeError: {e}\")" + ] + }, + { + "cell_type": "markdown", + "id": "a5b6c7d8", + "metadata": {}, + "source": [ + "### Pattern 2 — `eqx.filter_jit` (recommended)" + ] + }, + { + "cell_type": "code", + "execution_count": 39, + "id": "b6c7d8e9", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Controlled trajectory shape: (100, 2)\n" + ] + } + ], + "source": [ + "# apply_control is already @eqx.filter_jit\n", + "output_traj = ctrl.apply_control(control_guess, R_ctrl[-1])\n", + "print(f\"Controlled trajectory shape: {output_traj.shape}\")" + ] + }, + { + "cell_type": "markdown", + "id": "c7d8e9f0", + "metadata": {}, + "source": [ + "### Pattern 3 — `jax.jit` with explicit pytree argument" + ] + }, + { + "cell_type": "code", + "execution_count": 40, + "id": "d8e9f0a1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Controlled trajectory shape: (100, 2)\n" + ] + } + ], + "source": [ + "@jax.jit\n", + "def jit_apply_control(model, control_seq, res_state):\n", + " return model.apply_control(control_seq, res_state)\n", + "\n", + "output_traj = jit_apply_control(ctrl, control_guess, R_ctrl[-1])\n", + "print(f\"Controlled trajectory shape: {output_traj.shape}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/mkdocs.yml b/mkdocs.yml index 5c4f102..fd74ca1 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -116,6 +116,7 @@ nav: - Control: examples/control.ipynb - Data Library: examples/data_library.ipynb - Kuramoto-Sivashinsky: examples/ks.ipynb + - JAX JIT Compatibility: examples/jit_compatibility.ipynb - API Reference: - Forecaster: api/forecaster.md - Classifier: api/classifier.md diff --git a/src/orc/classifier/base.py b/src/orc/classifier/base.py index 868b87d..6090bd4 100644 --- a/src/orc/classifier/base.py +++ b/src/orc/classifier/base.py @@ -59,13 +59,13 @@ class RCClassifierBase(eqx.Module, ABC): driver: DriverBase readout: ReadoutBase embedding: EmbedBase - in_dim: int - out_dim: int - res_dim: int - n_classes: int - state_repr: str = "final" - dtype: type = jnp.float64 - seed: int = 0 + in_dim: int = eqx.field(static=True) + out_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) + n_classes: int = eqx.field(static=True) + state_repr: str = eqx.field(default="final", static=True) + dtype: type = eqx.field(default=jnp.float64, static=True) + seed: int = eqx.field(default=0, static=True) def __init__( self, diff --git a/src/orc/classifier/models.py b/src/orc/classifier/models.py index a9cf56c..0b07911 100644 --- a/src/orc/classifier/models.py +++ b/src/orc/classifier/models.py @@ -1,5 +1,6 @@ """Discrete ESN classifier implementation.""" +import equinox as eqx import jax import jax.numpy as jnp @@ -40,8 +41,8 @@ class ESNClassifier(RCClassifierBase): Replace embedding layer. """ - res_dim: int - data_dim: int + res_dim: int = eqx.field(static=True) + data_dim: int = eqx.field(static=True) def __init__( self, diff --git a/src/orc/control/base.py b/src/orc/control/base.py index 6931e7f..f7df458 100644 --- a/src/orc/control/base.py +++ b/src/orc/control/base.py @@ -68,15 +68,15 @@ class RCControllerBase(eqx.Module, ABC): driver: DriverBase readout: ReadoutBase embedding: EmbedBase - in_dim: int - control_dim: int - out_dim: int - res_dim: int - dtype: type = jnp.float64 + in_dim: int = eqx.field(static=True) + control_dim: int = eqx.field(static=True) + out_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) + dtype: type = eqx.field(default=jnp.float64, static=True) alpha_1: float = 100 alpha_2: float = 1 alpha_3: float = 5 - seed: int = 0 + seed: int = eqx.field(default=0, static=True) def __init__( self, diff --git a/src/orc/control/models.py b/src/orc/control/models.py index 505757b..ecaa195 100644 --- a/src/orc/control/models.py +++ b/src/orc/control/models.py @@ -1,5 +1,6 @@ """Discrete ESN controller implementation.""" +import equinox as eqx import jax import jax.numpy as jnp @@ -46,9 +47,9 @@ class ESNController(RCControllerBase): Replace embedding layer. """ - res_dim: int - data_dim: int - control_dim: int + res_dim: int = eqx.field(static=True) + data_dim: int = eqx.field(static=True) + control_dim: int = eqx.field(static=True) def __init__( self, diff --git a/src/orc/drivers.py b/src/orc/drivers.py index 1af0da0..276ce02 100644 --- a/src/orc/drivers.py +++ b/src/orc/drivers.py @@ -36,9 +36,9 @@ class DriverBase(eqx.Module, ABC): Create a zero initial reservoir state with appropriate shape. """ - res_dim: int - dtype: type - chunks: int = 0 + res_dim: int = eqx.field(static=True) + dtype: type = eqx.field(static=True) + chunks: int = eqx.field(default=0, static=True) def __init__(self, res_dim, dtype=jnp.float64): """Ensure reservoir dim and dtype are correct type.""" @@ -165,15 +165,15 @@ class ParallelESNDriver(DriverBase): Batched or single update to reservoir state. """ - res_dim: int + res_dim: int = eqx.field(static=True) leak: float spectral_radius: float density: float bias: float - dtype: type + dtype: type = eqx.field(static=True) wr: sparse.BCOO - chunks: int - mode: str + chunks: int = eqx.field(static=True) + mode: str = eqx.field(static=True) time_const: float def __init__( @@ -491,14 +491,14 @@ class ParallelTaylorDriver(DriverBase): Batched or single update to reservoir state. """ - n_terms: int - res_dim: int + n_terms: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) spectral_radius: float density: float bias: float - dtype: type + dtype: type = eqx.field(static=True) wr: sparse.BCOO - chunks: int + chunks: int = eqx.field(static=True) def __init__( self, @@ -854,7 +854,7 @@ class GRUDriver(DriverBase): """ gru: eqx.nn.GRUCell - chunks: int = 0 + chunks: int = eqx.field(default=0, static=True) def __init__(self, res_dim, input_rescaling=15.0, *, seed): """Initialize GRU-based reservoir driver. diff --git a/src/orc/embeddings.py b/src/orc/embeddings.py index bd2ba45..e421f40 100644 --- a/src/orc/embeddings.py +++ b/src/orc/embeddings.py @@ -31,10 +31,10 @@ class EmbedBase(eqx.Module, ABC): Embed multiple inputs into reservoir dimension. """ - in_dim: int - res_dim: int - dtype: type - chunks: int = 0 + in_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) + dtype: type = eqx.field(static=True) + chunks: int = eqx.field(default=0, static=True) def __init__(self, in_dim, res_dim, dtype=jnp.float64): """Ensure in dim, res dim, and dtype are correct type.""" @@ -148,15 +148,15 @@ class ParallelLinearEmbedding(EmbedBase): Embed single input state to reservoir dimension. """ - in_dim: int - res_dim: int + in_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) scaling: float win: Array - dtype: type - chunks: int - locality: int - chunk_size: int - periodic: bool + dtype: type = eqx.field(static=True) + chunks: int = eqx.field(static=True) + locality: int = eqx.field(static=True) + chunk_size: int = eqx.field(static=True) + periodic: bool = eqx.field(static=True) def __init__( self, @@ -387,12 +387,12 @@ class EnsembleLinearEmbedding(EmbedBase): Embed single input state to reservoir dimension. """ - in_dim: int - res_dim: int + in_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) scaling: float win: Array - dtype: type - chunks: int + dtype: type = eqx.field(static=True) + chunks: int = eqx.field(static=True) def __init__( self, diff --git a/src/orc/forecaster/base.py b/src/orc/forecaster/base.py index 2b1f62b..f59e8a2 100644 --- a/src/orc/forecaster/base.py +++ b/src/orc/forecaster/base.py @@ -58,12 +58,12 @@ class RCForecasterBase(eqx.Module, ABC): driver: DriverBase readout: ReadoutBase embedding: EmbedBase - in_dim: int - out_dim: int - res_dim: int - chunks: int = 0 - dtype: type = jnp.float64 - seed: int = 0 + in_dim: int = eqx.field(static=True) + out_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) + chunks: int = eqx.field(default=0, static=True) + dtype: type = eqx.field(default=jnp.float64, static=True) + seed: int = eqx.field(default=0, static=True) def __init__( self, diff --git a/src/orc/forecaster/models.py b/src/orc/forecaster/models.py index eb612f1..6b7501d 100644 --- a/src/orc/forecaster/models.py +++ b/src/orc/forecaster/models.py @@ -1,6 +1,7 @@ """Discrete and continuous ESN implementations with standard driver.""" import diffrax +import equinox as eqx import jax import jax.numpy as jnp @@ -45,8 +46,8 @@ class ESNForecaster(RCForecasterBase): Replace embedding layer. """ - res_dim: int - data_dim: int + res_dim: int = eqx.field(static=True) + data_dim: int = eqx.field(static=True) def __init__( self, @@ -180,8 +181,8 @@ class CESNForecaster(CRCForecasterBase): Replace embedding layer. """ - res_dim: int - data_dim: int + res_dim: int = eqx.field(static=True) + data_dim: int = eqx.field(static=True) def __init__( self, @@ -330,8 +331,8 @@ class EnsembleESNForecaster(RCForecasterBase): Replace embedding layer. """ - res_dim: int - data_dim: int + res_dim: int = eqx.field(static=True) + data_dim: int = eqx.field(static=True) def __init__( self, diff --git a/src/orc/readouts.py b/src/orc/readouts.py index cc3e280..c1da2ed 100644 --- a/src/orc/readouts.py +++ b/src/orc/readouts.py @@ -39,10 +39,10 @@ class ReadoutBase(eqx.Module, ABC): Return a new readout with updated output weights. """ - out_dim: int - res_dim: int - chunks: int = 0 - dtype: type = jnp.float64 + out_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) + chunks: int = eqx.field(default=0, static=True) + dtype: type = eqx.field(default=jnp.float64, static=True) def __init__(self, out_dim, res_dim, dtype=jnp.float64): """Ensure in dim, res dim, and dtype are correct type.""" @@ -200,11 +200,11 @@ class ParallelLinearReadout(ReadoutBase): Map from reservoir state to output state. """ - out_dim: int - res_dim: int + out_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) wout: Array - chunks: int - dtype: type + chunks: int = eqx.field(static=True) + dtype: type = eqx.field(static=True) def __init__( self, @@ -415,12 +415,12 @@ class ParallelNonlinearReadout(ReadoutBase): Map from reservoir state to output state, handles batch and single outputs. """ - out_dim: int - res_dim: int + out_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) wout: Array - chunks: int + chunks: int = eqx.field(static=True) nonlin_list: list - dtype: type + dtype: type = eqx.field(static=True) def __init__( self, @@ -734,11 +734,11 @@ class ParallelQuadraticReadout(ParallelNonlinearReadout): handles batch and single outputs. """ - out_dim: int - res_dim: int + out_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) wout: Array - chunks: int - dtype: type + chunks: int = eqx.field(static=True) + dtype: type = eqx.field(static=True) def __init__( self, @@ -851,11 +851,11 @@ class EnsembleLinearReadout(ReadoutBase): Map from reservoir state to output state. """ - out_dim: int - res_dim: int + out_dim: int = eqx.field(static=True) + res_dim: int = eqx.field(static=True) wout: Array - chunks: int - dtype: type + chunks: int = eqx.field(static=True) + dtype: type = eqx.field(static=True) def __init__( self, From 5b81c1d235a595ccd3de6647722dac02ad76cec9 Mon Sep 17 00:00:00 2001 From: Jan-Williams Date: Wed, 20 May 2026 09:16:18 -0700 Subject: [PATCH 2/2] Stripped jit compatibility notebook --- examples/jit_compatibility.ipynb | 217 ++++++++----------------------- 1 file changed, 55 insertions(+), 162 deletions(-) diff --git a/examples/jit_compatibility.ipynb b/examples/jit_compatibility.ipynb index 8d575f9..beba2db 100644 --- a/examples/jit_compatibility.ipynb +++ b/examples/jit_compatibility.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "a1b2c3d4", + "id": "0", "metadata": {}, "source": [ "# JAX JIT Compatibility\n", @@ -29,8 +29,8 @@ }, { "cell_type": "code", - "execution_count": 27, - "id": "b2c3d4e5", + "execution_count": null, + "id": "1", "metadata": {}, "outputs": [], "source": [ @@ -47,7 +47,7 @@ }, { "cell_type": "markdown", - "id": "c3d4e5f6", + "id": "2", "metadata": {}, "source": [ "## 1. Forecasting\n", @@ -57,18 +57,10 @@ }, { "cell_type": "code", - "execution_count": 28, - "id": "d4e5f6a7", + "execution_count": null, + "id": "3", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Reservoir states shape: (2400, 1, 500)\n" - ] - } - ], + "outputs": [], "source": [ "U, t = orc.data.lorenz63(tN=30, dt=0.01)\n", "split = int(0.8 * len(U))\n", @@ -82,7 +74,7 @@ }, { "cell_type": "markdown", - "id": "e5f6a7b8", + "id": "4", "metadata": {}, "source": [ "### Pattern 1 — closed-over `jax.jit` (does **not** work)\n", @@ -92,18 +84,10 @@ }, { "cell_type": "code", - "execution_count": 29, - "id": "f6a7b8c9", + "execution_count": null, + "id": "5", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "TypeError: unhashable type: 'BCOO'\n" - ] - } - ], + "outputs": [], "source": [ "try:\n", " U_pred = jax.jit(esn.forecast)(fcast_len=fcast_len, res_state=R[-1])\n", @@ -113,7 +97,7 @@ }, { "cell_type": "markdown", - "id": "a7b8c9d0", + "id": "6", "metadata": {}, "source": [ "### Pattern 2 — `eqx.filter_jit` (recommended)\n", @@ -125,19 +109,10 @@ }, { "cell_type": "code", - "execution_count": 30, - "id": "b8c9d0e1", + "execution_count": null, + "id": "7", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Forecast shape: (600, 3)\n", - "Forecast shape (explicit filter_jit): (600, 3)\n" - ] - } - ], + "outputs": [], "source": [ "# Calling the method directly compiles on first use (already @eqx.filter_jit)\n", "U_pred = esn.forecast(fcast_len=fcast_len, res_state=R[-1])\n", @@ -150,7 +125,7 @@ }, { "cell_type": "markdown", - "id": "c9d0e1f2", + "id": "8", "metadata": {}, "source": [ "### Pattern 3 — `jax.jit` with explicit pytree argument\n", @@ -160,18 +135,10 @@ }, { "cell_type": "code", - "execution_count": 31, - "id": "d0e1f2a3", + "execution_count": null, + "id": "9", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Forecast shape: (600, 3)\n" - ] - } - ], + "outputs": [], "source": [ "@functools.partial(jax.jit, static_argnums=(1,)) # fcast_len must be static\n", "def jit_forecast(model, fcast_len, res_state):\n", @@ -183,7 +150,7 @@ }, { "cell_type": "markdown", - "id": "e1f2a3b4", + "id": "10", "metadata": {}, "source": [ "### Training with `jax.jit`\n", @@ -193,19 +160,10 @@ }, { "cell_type": "code", - "execution_count": 32, - "id": "f2a3b4c5", + "execution_count": null, + "id": "11", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Training with jax.jit (defaults): OK\n", - "Training with jax.jit (explicit spinup): OK\n" - ] - } - ], + "outputs": [], "source": [ "esn_fresh = orc.forecaster.ESNForecaster(data_dim=3, res_dim=500, seed=1)\n", "\n", @@ -224,7 +182,7 @@ }, { "cell_type": "markdown", - "id": "a3b4c5d6", + "id": "12", "metadata": {}, "source": [ "## 2. Classification\n", @@ -234,18 +192,10 @@ }, { "cell_type": "code", - "execution_count": 33, - "id": "b4c5d6e7", + "execution_count": null, + "id": "13", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Classifier trained.\n" - ] - } - ], + "outputs": [], "source": [ "n_classes, data_dim, seq_len = 3, 3, 200\n", "key = jax.random.PRNGKey(42)\n", @@ -272,7 +222,7 @@ }, { "cell_type": "markdown", - "id": "c5d6e7f8", + "id": "14", "metadata": {}, "source": [ "### Pattern 1 — closed-over `jax.jit` (does **not** work)" @@ -280,18 +230,10 @@ }, { "cell_type": "code", - "execution_count": 34, - "id": "d6e7f8a9", + "execution_count": null, + "id": "15", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "TypeError: unhashable type: 'BCOO'\n" - ] - } - ], + "outputs": [], "source": [ "try:\n", " probs = jax.jit(clf.classify)(test_seqs[0])\n", @@ -301,7 +243,7 @@ }, { "cell_type": "markdown", - "id": "e7f8a9b0", + "id": "16", "metadata": {}, "source": [ "### Pattern 2 — `eqx.filter_jit` (recommended)" @@ -309,19 +251,10 @@ }, { "cell_type": "code", - "execution_count": 35, - "id": "f8a9b0c1", + "execution_count": null, + "id": "17", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Predicted class: 0, True class: 0\n", - "Test accuracy: 100.0%\n" - ] - } - ], + "outputs": [], "source": [ "# Single sample\n", "probs = clf.classify(test_seqs[0])\n", @@ -335,7 +268,7 @@ }, { "cell_type": "markdown", - "id": "a9b0c1d2", + "id": "18", "metadata": {}, "source": [ "### Pattern 3 — `jax.jit` with explicit pytree argument\n", @@ -345,18 +278,10 @@ }, { "cell_type": "code", - "execution_count": 36, - "id": "b0c1d2e3", + "execution_count": null, + "id": "19", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Predicted class: 0, True class: 0\n" - ] - } - ], + "outputs": [], "source": [ "@jax.jit\n", "def jit_classify(model, in_seq):\n", @@ -368,7 +293,7 @@ }, { "cell_type": "markdown", - "id": "c1d2e3f4", + "id": "20", "metadata": {}, "source": [ "## 3. Control\n", @@ -378,18 +303,10 @@ }, { "cell_type": "code", - "execution_count": 37, - "id": "d2e3f4a5", + "execution_count": null, + "id": "21", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Controller trained. Reservoir states shape: (599, 500)\n" - ] - } - ], + "outputs": [], "source": [ "# Simple stable 2D linear plant\n", "A_plant = jnp.array([[0.95, 0.1], [-0.1, 0.95]])\n", @@ -414,7 +331,7 @@ }, { "cell_type": "markdown", - "id": "e3f4a5b6", + "id": "22", "metadata": {}, "source": [ "### Pattern 1 — closed-over `jax.jit` (does **not** work)" @@ -422,18 +339,10 @@ }, { "cell_type": "code", - "execution_count": 38, - "id": "f4a5b6c7", + "execution_count": null, + "id": "23", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "TypeError: unhashable type: 'BCOO'\n" - ] - } - ], + "outputs": [], "source": [ "control_guess = jnp.zeros((100, 1))\n", "\n", @@ -445,7 +354,7 @@ }, { "cell_type": "markdown", - "id": "a5b6c7d8", + "id": "24", "metadata": {}, "source": [ "### Pattern 2 — `eqx.filter_jit` (recommended)" @@ -453,18 +362,10 @@ }, { "cell_type": "code", - "execution_count": 39, - "id": "b6c7d8e9", + "execution_count": null, + "id": "25", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Controlled trajectory shape: (100, 2)\n" - ] - } - ], + "outputs": [], "source": [ "# apply_control is already @eqx.filter_jit\n", "output_traj = ctrl.apply_control(control_guess, R_ctrl[-1])\n", @@ -473,7 +374,7 @@ }, { "cell_type": "markdown", - "id": "c7d8e9f0", + "id": "26", "metadata": {}, "source": [ "### Pattern 3 — `jax.jit` with explicit pytree argument" @@ -481,18 +382,10 @@ }, { "cell_type": "code", - "execution_count": 40, - "id": "d8e9f0a1", + "execution_count": null, + "id": "27", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Controlled trajectory shape: (100, 2)\n" - ] - } - ], + "outputs": [], "source": [ "@jax.jit\n", "def jit_apply_control(model, control_seq, res_state):\n",