{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "title",
   "metadata": {},
   "source": [
    "# Neural Network Layers for Online Learning\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "overview",
   "metadata": {},
   "source": [
    "After choosing the parameter operations that participate in online learning, compose them into reusable network blocks. `braintrace.nn` provides layers whose trainable forwards already use ETP operators. This tutorial explains operation-based selection, surveys the layer families, and builds a recurrent model whose parameter-to-hidden paths remain visible to the compiler.\n",
    "\n",
    "See the complete [Neural Network Layers API](../apis/nn.rst) for constructor signatures and generated class pages.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "operation-selection",
   "metadata": {},
   "source": [
    "## Selection is operation-based\n",
    "\n",
    "A parameter participates in online learning only when the forward path uses an ETP operation. {class}`braintrace.nn.Linear`, for example, uses {func}`braintrace.matmul`; the corresponding `brainstate.nn.Linear` uses an ordinary JAX operation. The class namespace is convenient, but the operation in the traced graph is what the compiler recognizes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "linear-example",
   "metadata": {},
   "outputs": [],
   "source": [
    "import brainstate\n",
    "import jax.numpy as jnp\n",
    "\n",
    "import braintrace\n",
    "\n",
    "brainstate.random.seed(23)\n",
    "linear = braintrace.nn.Linear(3, 2)\n",
    "linear_output = linear(jnp.ones(3))\n",
    "print(\"Linear output shape:\", linear_output.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "layer-families",
   "metadata": {},
   "source": [
    "## Layer families\n",
    "\n",
    "| Family | Representative API classes | ETP operation |\n",
    "|---|---|---|\n",
    "| Linear maps | {class}`braintrace.nn.Linear`, {class}`braintrace.nn.GroupedLinear`, {class}`braintrace.nn.SparseLinear`, {class}`braintrace.nn.LoRA` | {func}`braintrace.matmul`, {func}`braintrace.sparse_matmul`, {func}`braintrace.lora_matmul` |\n",
    "| Embeddings | {class}`braintrace.nn.Embedding` | indexed ETP-aware weight access |\n",
    "| Convolutions | {class}`braintrace.nn.Conv1d`, {class}`braintrace.nn.Conv2d`, {class}`braintrace.nn.Conv3d` | {func}`braintrace.conv` |\n",
    "| Recurrent cells | {class}`braintrace.nn.GRUCell`, {class}`braintrace.nn.LSTMCell`, {class}`braintrace.nn.MiniGRU`, {class}`braintrace.nn.MiniLSTM` | ETP dense and element-wise operations |\n",
    "| Readouts | {class}`braintrace.nn.LeakyRateReadout` | ETP-aware projection |\n",
    "\n",
    "Activation, normalization, and pooling layers should be imported directly from `brainstate.nn`; compatibility forwarding through `braintrace.nn` is deprecated. The [Neural Network Layers API](../apis/nn.rst) is the authoritative list of BrainTrace-owned layers.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "model-composition",
   "metadata": {},
   "source": [
    "## Compose and compile a recurrent model\n",
    "\n",
    "The recurrent layer writes hidden state, while the final Linear layer maps that state to an observable output. {func}`braintrace.compile` discovers the ETP relationships and returns the online learner.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compile-model",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TinySequenceModel(brainstate.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.rnn = braintrace.nn.MiniGRU(in_size=1, out_size=4)\n",
    "        self.readout = braintrace.nn.Linear(4, 1)\n",
    "\n",
    "    def update(self, x):\n",
    "        return self.readout(self.rnn(x))\n",
    "\n",
    "\n",
    "model = TinySequenceModel()\n",
    "sample = jnp.ones(1)\n",
    "learner = braintrace.compile(\n",
    "    model,\n",
    "    braintrace.D_RTRL,\n",
    "    sample,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "compiled-loop",
   "metadata": {},
   "source": [
    "## Run repeated steps with a compiled transform\n",
    "\n",
    "Use a `brainstate.transform` loop for repeated model execution. State is carried automatically and the outputs are stacked along the leading time axis.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "run-sequence",
   "metadata": {},
   "outputs": [],
   "source": [
    "sequence = jnp.linspace(-1.0, 1.0, 6).reshape(6, 1)\n",
    "brainstate.nn.reset_all_states(model)\n",
    "learner.reset_state()\n",
    "outputs = learner.etrace_evolve(sequence, return_outputs=True)\n",
    "print(\"Sequence output shape:\", outputs.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "temporal-parameters",
   "metadata": {},
   "source": [
    "## Temporal and non-temporal parameters\n",
    "\n",
    "The recurrent weights influence hidden state and therefore need eligibility traces. The readout weight is still trainable, but its output does not feed a hidden state, so the compiler classifies it as non-temporal and computes its instantaneous gradient directly.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "relation-boundaries",
   "metadata": {},
   "source": [
    "## Relation boundaries\n",
    "\n",
    "Trace every parameter path to hidden state before composing custom blocks. A `weight -> weight -> hidden` path crosses two trainable ETP operations. The compiler stops at the downstream operation, so the upstream weight is not recorded as an independent ETP relation; recording both would double-count a contribution that per-operation rules cannot represent jointly.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "next-steps",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "Continue with [Hidden States for Online Learning](hidden_states.ipynb), then compile and compare single-layer and stacked models in the [RNN Compiler Walkthrough](rnn_compiler.ipynb). Keep the [Neural Network Layers API](../apis/nn.rst) open when selecting a concrete layer.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}