{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ca38f518",
   "metadata": {},
   "source": [
    "# Customizing Primitive Transforms\n",
    "\n",
    "This is the *next chapter* after **ETP Primitives Deep Dive**. There we saw how each ETP primitive marks a weight operation and how a custom primitive (`scaled_matmul`) is registered. Here we add **parameter transform hooks** — small elementwise functions applied to a weight *inside* the primitive — and explain the one design rule that keeps online learning exact in their presence. We close with the optional closed-form **fast path** and its transform-aware gate."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c23d77d0",
   "metadata": {},
   "source": [
    "## 1. Why parameter transforms, and the one rule\n",
    "\n",
    "Models often need a *reparametrized* weight rather than the raw stored parameter:\n",
    "\n",
    "- **masked / structured** weights: `weight_fn=lambda w: w * mask`\n",
    "- **standardized** weights: subtract mean / divide by std before use\n",
    "- **sign-constrained** weights: `weight_fn=lambda w: w ** 2` (non-negative), `jnp.abs`, `jax.nn.softplus`\n",
    "- **squashed** weights: `weight_fn=jnp.tanh`\n",
    "\n",
    "Rather than rewriting the parameter and losing the raw value (which the optimizer updates), each built-in ETP op takes an optional, shape-preserving `*_fn` hook. The op computes the forward pass on the *transformed* weight $V = f(W)$ while the eligibility trace and gradient stay attached to the **raw** weight $W$.\n",
    "\n",
    "**The load-bearing rule.** Let $f$ be the transform and $V = f(W)$. The transform Jacobian $f'(W)$ is applied in **exactly one place** — the primitive's `xy_to_dw` rule, via `jax.vjp` through the impl. Everything else (`yw_to_w`, `init_drtrl`, `init_pp`) is transform-free: those rules operate on $\\partial h / \\partial W_{\\text{raw}}$ and are correct as written. Because `xy_to_dw` builds its VJP *through the same impl that applies $f$*, the chain rule threads $f'$ automatically — no per-transform code is needed.\n",
    "\n",
    "$$\n",
    "\\frac{\\partial h}{\\partial W_{\\text{raw}}}\n",
    "= \\underbrace{\\frac{\\partial h}{\\partial V}}_{\\text{closed-form rule}}\n",
    "  \\cdot \\underbrace{f'(W_{\\text{raw}})}_{\\text{threaded by } \\texttt{jax.vjp} \\text{ in } \\texttt{xy\\_to\\_dw}}\n",
    "$$\n",
    "\n",
    "Units note: every wrapper splits a `brainunit.Quantity` into mantissa + unit and applies the transform to the **unitless mantissa**, recombining the unit afterwards."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6387391b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-30T04:09:38.724504Z",
     "iopub.status.busy": "2026-06-30T04:09:38.724361Z",
     "iopub.status.idle": "2026-06-30T04:09:41.868101Z",
     "shell.execute_reply": "2026-06-30T04:09:41.867212Z"
    }
   },
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import brainunit as u\n",
    "import brainstate\n",
    "\n",
    "import braintrace\n",
    "\n",
    "brainstate.environ.set(precision=64)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40ccb969",
   "metadata": {},
   "source": [
    "## 2. Transforms on the built-in ops\n",
    "\n",
    "Each ETP op exposes the hooks that make sense for it. The names differ where the op has more than one trainable factor:\n",
    "\n",
    "| Op | Transform hook(s) |\n",
    "|---|---|\n",
    "| `braintrace.matmul(x, weight, bias=None, *, weight_fn=None, bias_fn=None)` | `weight_fn`, `bias_fn` |\n",
    "| `braintrace.element_wise(weight, *, weight_fn=None)` | `weight_fn` |\n",
    "| `braintrace.conv(x, kernel, bias=None, *, ..., kernel_fn=None, bias_fn=None)` | `kernel_fn` (note: *not* `weight_fn`), `bias_fn` |\n",
    "| `braintrace.sparse_matmul(x, weight, *, sparse_mat, bias=None, weight_fn=None, bias_fn=None)` | `weight_fn`, `bias_fn` |\n",
    "| `braintrace.lora_matmul(x, B, A, *, alpha=1.0, bias=None, b_fn=None, a_fn=None, bias_fn=None)` | per-factor `b_fn`, `a_fn`, plus `bias_fn` |\n",
    "\n",
    "All hooks default to `None` (identity), so omitting them reproduces the untransformed op bit-for-bit. The forward pass applies the transform:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fc2ffd41",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-30T04:09:41.870171Z",
     "iopub.status.busy": "2026-06-30T04:09:41.869945Z",
     "iopub.status.idle": "2026-06-30T04:10:04.109078Z",
     "shell.execute_reply": "2026-06-30T04:10:04.107531Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "matmul forward matches x @ tanh(w) + |b| : True\n"
     ]
    }
   ],
   "source": [
    "# matmul with weight_fn + bias_fn: forward uses the transformed weight/bias.\n",
    "x = brainstate.random.randn(4, 3)\n",
    "w = brainstate.random.randn(3, 5)\n",
    "b = brainstate.random.randn(5)\n",
    "\n",
    "y = braintrace.matmul(x, w, bias=b, weight_fn=jnp.tanh, bias_fn=jnp.abs)\n",
    "y_ref = x @ jnp.tanh(w) + jnp.abs(b)\n",
    "print(\"matmul forward matches x @ tanh(w) + |b| :\", bool(jnp.allclose(y, y_ref)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "04f46b2f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-30T04:10:04.113648Z",
     "iopub.status.busy": "2026-06-30T04:10:04.113360Z",
     "iopub.status.idle": "2026-06-30T04:10:17.729431Z",
     "shell.execute_reply": "2026-06-30T04:10:17.727833Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "element_wise matches tanh(w): True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "conv forward matches conv(x, kernel**2): True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lora forward matches alpha * x @ B @ tanh(A): True\n"
     ]
    }
   ],
   "source": [
    "# element_wise: the single weight_fn squashes the parameter in place.\n",
    "wg = brainstate.random.randn(5)\n",
    "print(\"element_wise matches tanh(w):\",\n",
    "      bool(jnp.allclose(braintrace.element_wise(wg, weight_fn=jnp.tanh), jnp.tanh(wg))))\n",
    "\n",
    "# conv uses kernel_fn (not weight_fn) for the kernel.\n",
    "xc = brainstate.random.randn(8, 3, 16)        # NCH (JAX default layout)\n",
    "k = brainstate.random.randn(4, 3, 5)          # OIH\n",
    "yc = braintrace.conv(xc, k, strides=(1,), padding='SAME', kernel_fn=lambda ww: ww ** 2)\n",
    "yc_ref = jax.lax.conv_general_dilated(xc, k ** 2, window_strides=(1,), padding='SAME')\n",
    "print(\"conv forward matches conv(x, kernel**2):\", bool(jnp.allclose(yc, yc_ref)))\n",
    "\n",
    "# lora has per-factor b_fn / a_fn.\n",
    "xl = brainstate.random.randn(16, 8)\n",
    "B = brainstate.random.randn(8, 2)\n",
    "A = brainstate.random.randn(2, 4)\n",
    "yl = braintrace.lora_matmul(xl, B, A, alpha=0.5, a_fn=jnp.tanh)\n",
    "print(\"lora forward matches alpha * x @ B @ tanh(A):\",\n",
    "      bool(jnp.allclose(yl, 0.5 * (xl @ B @ jnp.tanh(A)))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9c2d84de",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-30T04:10:17.732106Z",
     "iopub.status.busy": "2026-06-30T04:10:17.731847Z",
     "iopub.status.idle": "2026-06-30T04:10:17.868021Z",
     "shell.execute_reply": "2026-06-30T04:10:17.865725Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "unit preserved (V * S = A): A\n"
     ]
    }
   ],
   "source": [
    "# Units act on the mantissa: the transform sees the unitless value, the unit is\n",
    "# split off before and recombined after.\n",
    "x_q = jnp.ones((4, 3)) * u.volt\n",
    "w_q = jnp.ones((3, 5)) * u.siemens\n",
    "y_q = braintrace.matmul(x_q, w_q, weight_fn=lambda ww: ww ** 2)\n",
    "print(\"unit preserved (V * S = A):\", u.get_unit(y_q))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "582653e6",
   "metadata": {},
   "source": [
    "## 3. Correctness check: a transformed weight is handled *exactly*\n",
    "\n",
    "A transform is not an approximation. Because `xy_to_dw` threads $f'$ via `jax.vjp`, an **exact** online algorithm such as `D_RTRL` must still reproduce the BPTT gradient element-wise. We verify this with the gradient oracle in `braintrace._algorithm.oracle`:\n",
    "\n",
    "- `bptt_param_gradients(factory, inputs)` — exact backprop-through-time over the sequence sum-of-squares loss.\n",
    "- `online_param_gradients(factory, inputs, algo_factory=...)` — the same total gradient via the online algorithm's multi-step path.\n",
    "- `assert_param_gradients_close(online, bptt, atol=...)` — element-wise comparison.\n",
    "\n",
    "We build a tiny RNN cell whose recurrent map routes through `braintrace.matmul(..., weight_fn=jnp.tanh)` and assert `D_RTRL == BPTT`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1791e74c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-30T04:10:17.872301Z",
     "iopub.status.busy": "2026-06-30T04:10:17.871984Z",
     "iopub.status.idle": "2026-06-30T04:10:31.712950Z",
     "shell.execute_reply": "2026-06-30T04:10:31.711941Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "D_RTRL == BPTT with weight_fn=tanh  ->  transform handled exactly\n"
     ]
    }
   ],
   "source": [
    "from braintrace._algorithm.oracle import (\n",
    "    bptt_param_gradients,\n",
    "    online_param_gradients,\n",
    "    assert_param_gradients_close,\n",
    ")\n",
    "\n",
    "\n",
    "class TanhWeightRNN(brainstate.nn.Module):\n",
    "    \"\"\"A 1-state-group RNN whose weight is squashed by tanh inside the ETP op.\"\"\"\n",
    "\n",
    "    def __init__(self, in_dim=2, hid_dim=4):\n",
    "        super().__init__()\n",
    "        self.in_dim = in_dim\n",
    "        self.hid_dim = hid_dim\n",
    "        # One weight maps [x ; h] -> h, marked with braintrace.matmul.\n",
    "        self.W = brainstate.ParamState(\n",
    "            brainstate.random.randn(in_dim + hid_dim, hid_dim) * 0.2\n",
    "        )\n",
    "\n",
    "    def init_state(self, batch_size=None, **kwargs):\n",
    "        size = (self.hid_dim,) if batch_size is None else (batch_size, self.hid_dim)\n",
    "        self.h = brainstate.HiddenState(jnp.zeros(size))\n",
    "\n",
    "    def update(self, x):\n",
    "        xh = jnp.concatenate([x.reshape(1, -1), self.h.value], axis=-1)\n",
    "        # weight_fn=tanh: forward uses tanh(W); the trace stays attached to raw W.\n",
    "        self.h.value = jnp.tanh(braintrace.matmul(xh, self.W.value, weight_fn=jnp.tanh))\n",
    "        return self.h.value\n",
    "\n",
    "\n",
    "def factory():\n",
    "    brainstate.random.seed(0)\n",
    "    return TanhWeightRNN()\n",
    "\n",
    "\n",
    "brainstate.random.seed(1)\n",
    "inputs = brainstate.random.randn(6, 2)  # 6 time steps, in_dim=2\n",
    "\n",
    "bptt = bptt_param_gradients(factory, inputs)\n",
    "online = online_param_gradients(\n",
    "    factory, inputs,\n",
    "    algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step'),\n",
    ")\n",
    "assert_param_gradients_close(online, bptt, atol=1e-4)\n",
    "print(\"D_RTRL == BPTT with weight_fn=tanh  ->  transform handled exactly\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85c73db3",
   "metadata": {},
   "source": [
    "## 4. Adding transforms to a custom primitive\n",
    "\n",
    "We now extend the `scaled_matmul` example from the previous tutorial — $y = \\text{scale} \\cdot (x\\,@\\,W) \\;(+ b)$ — with transform hooks. **The four ETP rules are unchanged from the untransformed version.** Only the *impl* (and the `xy_to_dw` forward it differentiates) need to apply `weight_fn` / `bias_fn`; because `xy_to_dw` builds its VJP through that same forward, $f'$ is threaded for free and `D_RTRL` stays exact."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "30ae0189",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-30T04:10:31.715494Z",
     "iopub.status.busy": "2026-06-30T04:10:31.715271Z",
     "iopub.status.idle": "2026-06-30T04:10:31.719393Z",
     "shell.execute_reply": "2026-06-30T04:10:31.718301Z"
    }
   },
   "outputs": [],
   "source": [
    "from braintrace import register_primitive\n",
    "\n",
    "\n",
    "# Step 1: impl applies the transforms (this is the ONLY place the forward changes).\n",
    "def _scaled_matmul_impl(*args, scale=1.0, has_bias=False, weight_fn=None, bias_fn=None):\n",
    "    x, w = args[0], args[1]\n",
    "    if weight_fn is not None:\n",
    "        w = weight_fn(w)\n",
    "    y = scale * (x @ w)\n",
    "    if has_bias:\n",
    "        b = args[2]\n",
    "        if bias_fn is not None:\n",
    "            b = bias_fn(b)\n",
    "        y = y + b\n",
    "    return y\n",
    "\n",
    "\n",
    "def _scaled_trainable_invars(params):\n",
    "    base = {'weight': 1}\n",
    "    if params.get('has_bias', False):\n",
    "        base['bias'] = 2\n",
    "    return base\n",
    "\n",
    "\n",
    "scaled_mm_p = register_primitive(\n",
    "    'etp_scaled_mm_transforms',\n",
    "    _scaled_matmul_impl,\n",
    "    batched=True,\n",
    "    trainable_invars_fn=_scaled_trainable_invars,\n",
    "    x_invar_index=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "508ef4c6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-30T04:10:31.721023Z",
     "iopub.status.busy": "2026-06-30T04:10:31.720890Z",
     "iopub.status.idle": "2026-06-30T04:10:31.726327Z",
     "shell.execute_reply": "2026-06-30T04:10:31.725521Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "custom primitive registered with transform-aware xy_to_dw\n"
     ]
    }
   ],
   "source": [
    "# Step 2: the four ETP rules. Note: yw_to_w / init_* are TRANSFORM-FREE.\n",
    "# Only xy_to_dw differentiates the impl (which applies f), so f' enters here alone.\n",
    "\n",
    "def _scaled_yw_to_w(hidden_dim, trace, *, scale=1.0, has_bias=False,\n",
    "                    weight_fn=None, bias_fn=None):\n",
    "    # y -> W chain link for y = scale * x @ W: scale on the matching out column.\n",
    "    # No f' here by design.\n",
    "    out = {'weight': trace['weight'] * jnp.expand_dims(hidden_dim, axis=-2) * scale}\n",
    "    if has_bias:\n",
    "        out['bias'] = trace['bias'] * hidden_dim\n",
    "    return out\n",
    "\n",
    "\n",
    "def _scaled_xy_to_dw(x, hidden_dim, weights, *, scale=1.0, has_bias=False,\n",
    "                     weight_fn=None, bias_fn=None):\n",
    "    # Single fused VJP THROUGH the transform -> gradient w.r.t. the RAW weight,\n",
    "    # with f' threaded automatically by jax.vjp.\n",
    "    def _fwd(w_dict):\n",
    "        w = w_dict['weight']\n",
    "        if weight_fn is not None:\n",
    "            w = weight_fn(w)\n",
    "        y = scale * (x @ w)\n",
    "        if has_bias:\n",
    "            b = w_dict['bias']\n",
    "            if bias_fn is not None:\n",
    "                b = bias_fn(b)\n",
    "            y = y + b\n",
    "        return u.get_mantissa(y)\n",
    "\n",
    "    _, vjp_fn = jax.vjp(_fwd, weights)\n",
    "    return jax.tree.map(u.get_mantissa, vjp_fn(hidden_dim)[0])\n",
    "\n",
    "\n",
    "def _scaled_init_drtrl(x_var, y_var, weight_vars, num_hidden_state):\n",
    "    batch = x_var.aval.shape[0]\n",
    "    out = {'weight': jnp.zeros((batch, *weight_vars['weight'].aval.shape, num_hidden_state))}\n",
    "    if 'bias' in weight_vars:\n",
    "        out['bias'] = jnp.zeros((batch, *weight_vars['bias'].aval.shape, num_hidden_state))\n",
    "    return out\n",
    "\n",
    "\n",
    "def _scaled_init_pp(x_var, y_var, weight_vars, num_hidden_state):\n",
    "    return jnp.zeros((*y_var.aval.shape, num_hidden_state), dtype=y_var.aval.dtype)\n",
    "\n",
    "\n",
    "scaled_mm_p.register_etp_rules(\n",
    "    yw_to_w=_scaled_yw_to_w,\n",
    "    xy_to_dw=_scaled_xy_to_dw,\n",
    "    init_drtrl=_scaled_init_drtrl,\n",
    "    init_pp=_scaled_init_pp,\n",
    ")\n",
    "print(\"custom primitive registered with transform-aware xy_to_dw\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2dc77b4c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-30T04:10:31.727866Z",
     "iopub.status.busy": "2026-06-30T04:10:31.727727Z",
     "iopub.status.idle": "2026-06-30T04:10:38.044842Z",
     "shell.execute_reply": "2026-06-30T04:10:38.043871Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "forward matches 2 * (x @ tanh(w)): True\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "grad shape: (3, 5)\n"
     ]
    }
   ],
   "source": [
    "# Forward + grad work immediately (auto-derived JAX rules thread the transform).\n",
    "x = jnp.ones((4, 3))\n",
    "w = jnp.ones((3, 5))\n",
    "y = scaled_mm_p.bind(x, w, scale=2.0, has_bias=False, weight_fn=jnp.tanh, bias_fn=None)\n",
    "print(\"forward matches 2 * (x @ tanh(w)):\", bool(jnp.allclose(y, 2.0 * (x @ jnp.tanh(w)))))\n",
    "\n",
    "dw = jax.grad(\n",
    "    lambda w_: jnp.sum(\n",
    "        scaled_mm_p.bind(x, w_, scale=2.0, has_bias=False, weight_fn=jnp.tanh, bias_fn=None)\n",
    "    )\n",
    ")(w)\n",
    "print(\"grad shape:\", dw.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "92a7a949",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-30T04:10:38.047374Z",
     "iopub.status.busy": "2026-06-30T04:10:38.047106Z",
     "iopub.status.idle": "2026-06-30T04:10:38.753482Z",
     "shell.execute_reply": "2026-06-30T04:10:38.752496Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "custom scaled_mm: D_RTRL == BPTT with weight_fn=tanh  ->  rules need no change\n"
     ]
    }
   ],
   "source": [
    "# And the four unchanged ETP rules give an EXACT D_RTRL gradient with the\n",
    "# transform active, just like the built-in matmul did in Section 3.\n",
    "class ScaledTanhRNN(brainstate.nn.Module):\n",
    "    def __init__(self, in_dim=2, hid_dim=4):\n",
    "        super().__init__()\n",
    "        self.in_dim = in_dim\n",
    "        self.hid_dim = hid_dim\n",
    "        self.W = brainstate.ParamState(\n",
    "            brainstate.random.randn(in_dim + hid_dim, hid_dim) * 0.2\n",
    "        )\n",
    "\n",
    "    def init_state(self, batch_size=None, **kwargs):\n",
    "        size = (self.hid_dim,) if batch_size is None else (batch_size, self.hid_dim)\n",
    "        self.h = brainstate.HiddenState(jnp.zeros(size))\n",
    "\n",
    "    def update(self, x):\n",
    "        xh = jnp.concatenate([x.reshape(1, -1), self.h.value], axis=-1)\n",
    "        r = scaled_mm_p.bind(xh, self.W.value, scale=1.5, has_bias=False,\n",
    "                             weight_fn=jnp.tanh, bias_fn=None)\n",
    "        self.h.value = jnp.tanh(r)\n",
    "        return self.h.value\n",
    "\n",
    "\n",
    "def custom_factory():\n",
    "    brainstate.random.seed(0)\n",
    "    return ScaledTanhRNN()\n",
    "\n",
    "\n",
    "brainstate.random.seed(1)\n",
    "inp = brainstate.random.randn(6, 2)\n",
    "bptt_c = bptt_param_gradients(custom_factory, inp)\n",
    "online_c = online_param_gradients(\n",
    "    custom_factory, inp,\n",
    "    algo_factory=lambda m: braintrace.D_RTRL(m, vjp_method='multi-step'),\n",
    ")\n",
    "assert_param_gradients_close(online_c, bptt_c, atol=1e-4)\n",
    "print(\"custom scaled_mm: D_RTRL == BPTT with weight_fn=tanh  ->  rules need no change\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95f8f50a",
   "metadata": {},
   "source": [
    "## 5. (Advanced) The fast path and its transform-aware gate\n",
    "\n",
    "For the elementwise-`yw_to_w` primitives (`etp_mm` / `etp_mv` / `etp_elemwise`), a closed-form param-dim D-RTRL kernel bundle replaces the generic nested-`vmap` trace path with direct einsums. The bundle is a `FastPathRules(instant, recurrent, solve, applicable)` namedtuple, registered alongside the four rules via `register_etp_rules(fast_path=...)`, and looked up with `get_fast_path_rules(primitive)`.\n",
    "\n",
    "- `instant` — the instantaneous $\\operatorname{diag}(\\mathbf{D}_f^t) \\otimes \\mathbf{x}^t$ term.\n",
    "- `recurrent` — the $\\mathbf{D}^t \\boldsymbol{\\epsilon}^{t-1}$ contraction.\n",
    "- `solve` — the solve-time contraction of the learning signal with the trace.\n",
    "- `applicable(eqn_params)` — the **gate**.\n",
    "\n",
    "The closed-form kernels compute $\\partial h / \\partial V$ for the *transformed* weight $V = f(W)$ — they emit the bare outer product and **drop $f'$**. So whenever a transform hook is present, `applicable` returns `False`, and that relation falls back to the rule path (whose `xy_to_dw` applies $f'$ via `jax.vjp`, as in Sections 3-4). This is what keeps the fast path an optimization, never a correctness hazard.\n",
    "\n",
    "Note `sparse` / `conv` / `lora` have **no** fast path at all — they always use the rule path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1bf1f341",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-30T04:10:38.755192Z",
     "iopub.status.busy": "2026-06-30T04:10:38.755031Z",
     "iopub.status.idle": "2026-06-30T04:10:38.759114Z",
     "shell.execute_reply": "2026-06-30T04:10:38.758302Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "etp_mm has a fast path: True\n",
      "applicable (no transform)     : True\n",
      "applicable (weight_fn=tanh)    : False\n",
      "etp_elemwise fast path: True\n",
      "etp_sp_mm fast path is None: True\n",
      "etp_conv fast path is None: True\n",
      "etp_lora_mm fast path is None: True\n"
     ]
    }
   ],
   "source": [
    "from braintrace._op import (\n",
    "    get_fast_path_rules,\n",
    "    etp_mm_p, etp_elemwise_p,\n",
    "    etp_sp_mm_p, etp_conv_p, etp_lora_mm_p,\n",
    ")\n",
    "\n",
    "fp = get_fast_path_rules(etp_mm_p)\n",
    "assert fp is not None\n",
    "print(\"etp_mm has a fast path:\", fp is not None)\n",
    "\n",
    "# Gate ON when no transform; OFF the moment a transform hook is present.\n",
    "print(\"applicable (no transform)     :\", fp.applicable({'weight_fn': None, 'bias_fn': None}))\n",
    "print(\"applicable (weight_fn=tanh)    :\", fp.applicable({'weight_fn': jnp.tanh, 'bias_fn': None}))\n",
    "assert fp.applicable({'weight_fn': None, 'bias_fn': None}) is True\n",
    "assert fp.applicable({'weight_fn': jnp.tanh, 'bias_fn': None}) is False\n",
    "\n",
    "# elemwise also has one; sparse / conv / lora do not.\n",
    "print(\"etp_elemwise fast path:\", get_fast_path_rules(etp_elemwise_p) is not None)\n",
    "for p, name in [(etp_sp_mm_p, 'etp_sp_mm'), (etp_conv_p, 'etp_conv'), (etp_lora_mm_p, 'etp_lora_mm')]:\n",
    "    print(f\"{name} fast path is None:\", get_fast_path_rules(p) is None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35baaa47",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "- **Transform hooks** let a model use a reparametrized weight (`weight_fn` / `kernel_fn` / `b_fn` / `a_fn` / `bias_fn`) while the trace and gradient stay attached to the raw parameter.\n",
    "- **One rule:** the transform Jacobian $f'$ enters *only* in `xy_to_dw` (via `jax.vjp` through the impl). `yw_to_w` and the trace initializers stay transform-free and exact.\n",
    "- **Exactness:** `D_RTRL == BPTT` element-wise even with `weight_fn=tanh`, for both built-in and custom primitives — verified via `braintrace._algorithm.oracle`.\n",
    "- **Custom primitives** need no rule changes for transforms: apply the hook in the impl, and `xy_to_dw`'s VJP threads $f'$ for free.\n",
    "- **Fast path:** the closed-form `FastPathRules` bundle drops $f'$, so its `applicable` gate disables it under any transform; such relations fall back to the (correct) rule path. `sparse` / `conv` / `lora` have no fast path."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}