diff --git a/README.md b/README.md
index 2201c37..507079d 100644
--- a/README.md
+++ b/README.md
@@ -65,6 +65,9 @@ plt.show()
   <img src="imgs/readme_example_forecast.png" alt="ORC Logo"/>
 </div>
 
+## 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..beba2db
--- /dev/null
+++ b/examples/jit_compatibility.ipynb
@@ -0,0 +1,412 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "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": null,
+   "id": "1",
+   "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": "2",
+   "metadata": {},
+   "source": [
+    "## 1. Forecasting\n",
+    "\n",
+    "Train an `ESNForecaster` on the Lorenz-63 system."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "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": "4",
+   "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": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "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": "6",
+   "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": null,
+   "id": "7",
+   "metadata": {},
+   "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",
+    "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": "8",
+   "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": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "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": "10",
+   "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": null,
+   "id": "11",
+   "metadata": {},
+   "outputs": [],
+   "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": "12",
+   "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": null,
+   "id": "13",
+   "metadata": {},
+   "outputs": [],
+   "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": "14",
+   "metadata": {},
+   "source": [
+    "### Pattern 1 — closed-over `jax.jit` (does **not** work)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "15",
+   "metadata": {},
+   "outputs": [],
+   "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": "16",
+   "metadata": {},
+   "source": [
+    "### Pattern 2 — `eqx.filter_jit` (recommended)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "17",
+   "metadata": {},
+   "outputs": [],
+   "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": "18",
+   "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": null,
+   "id": "19",
+   "metadata": {},
+   "outputs": [],
+   "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": "20",
+   "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": null,
+   "id": "21",
+   "metadata": {},
+   "outputs": [],
+   "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": "22",
+   "metadata": {},
+   "source": [
+    "### Pattern 1 — closed-over `jax.jit` (does **not** work)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23",
+   "metadata": {},
+   "outputs": [],
+   "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": "24",
+   "metadata": {},
+   "source": [
+    "### Pattern 2 — `eqx.filter_jit` (recommended)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "25",
+   "metadata": {},
+   "outputs": [],
+   "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": "26",
+   "metadata": {},
+   "source": [
+    "### Pattern 3 — `jax.jit` with explicit pytree argument"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "27",
+   "metadata": {},
+   "outputs": [],
+   "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,