{ "cells": [ { "cell_type": "markdown", "id": "single-title", "metadata": {}, "source": [ "# RNN Compiler Walkthrough" ] }, { "cell_type": "markdown", "id": "single-intro", "metadata": {}, "source": [ "This tutorial follows the same compiler workflow from a single recurrent layer to a stacked RNN. We first establish how to read the compilation report and `ETraceGraph`, then compile a two-layer model and compare the structures reported for both models." ] }, { "cell_type": "markdown", "id": "single-model-heading", "metadata": {}, "source": [ "## Single-Layer RNN\n", "\n", "### Define and Compile the Model\n", "\n", "The model contains one `ValinaRNNCell` and one linear readout. The recurrent cell owns the temporal hidden state, while the readout maps that state to the final output." ] }, { "cell_type": "code", "execution_count": 1, "id": "single-imports", "metadata": {}, "outputs": [], "source": [ "import jax\n", "import jax.numpy as jnp\n", "import brainstate\n", "import braintrace" ] }, { "cell_type": "code", "execution_count": 2, "id": "single-compile", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "========================================================================================================================\n", "The hidden groups are:\n", "\n", " Group 0: [('rnn', 'h')]\n", "\n", "\n", "The weight parameters which are associated with the hidden states are:\n", "\n", " Weight 0: ('rnn', 'W', 'weight') is associated with hidden group 0\n", "\n", "\n", "The non-etrace weight parameters are:\n", "\n", " Weight 0: ('out', 'weight') (excluded: relation_excluded_non_temporal)\n", "\n", "\n", "Compiler diagnostics (warnings / errors):\n", "\n", " [warning] relation_excluded_non_temporal: ETP primitive etp_mv (weight=('out', 'weight')) has no connected hidden states. It will be treated as a non-temporal parameter.\n", "\n", "\n", "\n", "========================================================================================================================\n", "The hidden groups are:\n", "\n", " Group 0: [('rnn', 'h')]\n", "\n", "\n", "The weight parameters which are associated with the hidden states are:\n", "\n", " Weight 0: ('rnn', 'W', 'weight') is associated with hidden group 0\n", "\n", "\n", "The non-etrace weight parameters are:\n", "\n", " Weight 0: ('out', 'weight') (excluded: relation_excluded_non_temporal)\n", "\n", "\n", "\n" ] } ], "source": [ "class SingleLayerRNN(brainstate.nn.Module):\n", " def __init__(self, n_in, n_rec, n_out):\n", " super().__init__()\n", " self.rnn = braintrace.nn.ValinaRNNCell(n_in, n_rec)\n", " self.out = braintrace.nn.Linear(n_rec, n_out)\n", "\n", " def update(self, x):\n", " return self.out(self.rnn(x))\n", "\n", "\n", "model = SingleLayerRNN(10, 32, 5)\n", "\n", "# braintrace.compile initialises states, compiles the ETP graph, and returns a ready learner.\n", "# We compile for a single unbatched sample (no batch_size), so the hidden state is (32,) and\n", "# the recurrent op is the matrix-vector primitive etp_mv. verbose=2 prints full diagnostics.\n", "learner = braintrace.compile(model, braintrace.D_RTRL, jnp.zeros(10), verbose=2)\n", "learner.show_graph()" ] }, { "cell_type": "markdown", "id": "single-output", "metadata": {}, "source": [ "### Read the Compiler Output\n", "\n", "The compiler identifies one hidden group at `('rnn', 'h')`. The recurrent weight at `('rnn', 'W', 'weight')` reaches that hidden state through an ETP primitive, so the learner maintains an eligibility trace for it.\n", "\n", "The readout weight at `('out', 'weight')` is different. Its output does not feed a recurrent hidden state, so it is reported as `relation_excluded_non_temporal`. It remains trainable through the loss, but it does not require a temporal eligibility trace." ] }, { "cell_type": "markdown", "id": "single-report-heading", "metadata": {}, "source": [ "### Using `learner.report` -- the `CompilationReport`\n", "\n", "Every learner returned by `braintrace.compile` exposes a `CompilationReport` through `learner.report`. The report is the concise view of what compilation included, excluded, or diagnosed:\n", "\n", "- `report.counts` summarizes hidden groups, ETP weights, excluded weights, warnings, and errors.\n", "- `report.etrace_weights` lists parameter paths that participate in eligibility tracing.\n", "- `report.excluded_weights` pairs excluded parameter paths with their reasons.\n", "- `report.dynamic_states` records non-hidden dynamic states discovered during tracing.\n", "- `report.diagnostics` contains the complete `CompilationRecord` sequence." ] }, { "cell_type": "code", "execution_count": 3, "id": "single-report", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "========================================================================================================================\n", "The hidden groups are:\n", "\n", " Group 0: [('rnn', 'h')]\n", "\n", "\n", "The weight parameters which are associated with the hidden states are:\n", "\n", " Weight 0: ('rnn', 'W', 'weight') is associated with hidden group 0\n", "\n", "\n", "The non-etrace weight parameters are:\n", "\n", " Weight 0: ('out', 'weight') (excluded: relation_excluded_non_temporal)\n", "\n", "\n", "\n", "Counts: {'hidden_groups': 1, 'etrace_weights': 2, 'excluded_weights': 1, 'warnings': 1, 'errors': 0}\n", "ETrace weights: [(('rnn', 'W', 'weight'), [0]), (('rnn', 'W', 'weight'), [0])]\n", "Excluded weights: [(('out', 'weight'), 'relation_excluded_non_temporal')]\n" ] } ], "source": [ "# report.show(level) prints a structured summary at the requested verbosity.\n", "# level=1 shows hidden groups, etrace weights, and excluded weights.\n", "learner.report.show(1)\n", "\n", "# Programmatic access to the summary counts\n", "print(\"Counts:\", learner.report.counts)\n", "\n", "# Which weights participate in online learning?\n", "print(\"ETrace weights:\", learner.report.etrace_weights)\n", "\n", "# Which weights were excluded (e.g., non-temporal readouts)?\n", "print(\"Excluded weights:\", learner.report.excluded_weights)" ] }, { "cell_type": "markdown", "id": "single-graph-heading", "metadata": {}, "source": [ "### Understanding `ETraceGraph`\n", "\n", "The report summarizes decisions; `learner.graph` exposes their structural representation. Its central fields are:\n", "\n", "| Field | Meaning |\n", "|---|---|\n", "| `module_info` | Traced Jaxpr and model-state mappings |\n", "| `hidden_groups` | Hidden states updated as recurrent groups |\n", "| `hid_path_to_group` | Hidden-state path to group mapping |\n", "| `hidden_param_op_relations` | Parameter, ETP primitive, and hidden-group relations |\n", "| `hidden_perturb` | Perturbation structure used for hidden Jacobians |\n", "| `diagnostics` | Included and excluded compiler decisions |\n", "\n", "Inspecting the relation itself confirms which trainable path, primitive, and hidden group produced the summary above." ] }, { "cell_type": "code", "execution_count": 4, "id": "single-graph", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "=== Hidden Groups ===\n", " Group 0: 1 state(s), shape (32,)\n", " Paths: [('rnn', 'h')]\n", "\n", "=== Weight-Primitive-Hidden Relations ===\n", " Relation 0:\n", " Trainable paths: {'weight': ('rnn', 'W', 'weight'), 'bias': ('rnn', 'W', 'weight')}\n", " Primitive: etp_mv\n", " Hidden groups: [0]\n", "\n", "=== Perturbation ===\n", " Has perturbation: True\n" ] } ], "source": [ "graph = learner.graph\n", "\n", "print(\"=== Hidden Groups ===\")\n", "for g in graph.hidden_groups:\n", " print(f\" Group {g.index}: {g.num_state} state(s), shape {g.varshape}\")\n", " print(f\" Paths: {g.hidden_paths}\")\n", "\n", "print(\"\\n=== Weight-Primitive-Hidden Relations ===\")\n", "for i, r in enumerate(graph.hidden_param_op_relations):\n", " print(f\" Relation {i}:\")\n", " # ``trainable_paths`` is a dict {trainable key -> owning ParamState path};\n", " # a single primitive may own several (e.g. {weight, bias}).\n", " print(f\" Trainable paths: {r.trainable_paths}\")\n", " print(f\" Primitive: {r.primitive}\")\n", " print(f\" Hidden groups: {[g.index for g in r.hidden_groups]}\")\n", "\n", "print(f\"\\n=== Perturbation ===\")\n", "print(f\" Has perturbation: {graph.hidden_perturb is not None}\")" ] }, { "cell_type": "markdown", "id": "single-direct-heading", "metadata": {}, "source": [ "### Using `compile_etrace_graph` Directly\n", "\n", "`braintrace.compile(...)` is the standard entry point because it initializes states, compiles the graph, and returns a ready learner. For structural debugging or custom algorithm development, `braintrace.compile_etrace_graph(...)` returns the same `ETraceGraph` without constructing an algorithm wrapper." ] }, { "cell_type": "code", "execution_count": 5, "id": "single-direct", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Number of hidden groups: 1\n", "Number of relations: 1\n", "Has perturbation: True\n", "\n", "Graph fields:\n", " module_info\n", " hidden_groups\n", " hid_path_to_group\n", " hidden_param_op_relations\n", " hidden_perturb\n", " diagnostics\n" ] } ], "source": [ "model_direct = SingleLayerRNN(10, 32, 5)\n", "brainstate.nn.init_all_states(model_direct)\n", "\n", "graph_direct = braintrace.compile_etrace_graph(model_direct, jnp.zeros(10))\n", "\n", "print(f\"Number of hidden groups: {len(graph_direct.hidden_groups)}\")\n", "print(f\"Number of relations: {len(graph_direct.hidden_param_op_relations)}\")\n", "print(f\"Has perturbation: {graph_direct.hidden_perturb is not None}\")\n", "\n", "print(\"\\nGraph fields:\")\n", "for key in graph_direct.dict().keys():\n", " print(f\" {key}\")" ] }, { "cell_type": "markdown", "id": "two-model-heading", "metadata": {}, "source": [ "## Two-Layer RNN\n", "\n", "The stacked model uses the same input and hidden width but introduces a second recurrent state. It uses `GRUCell` layers to expose multiple gate relations as well as multiple hidden groups; consequently, the final comparison is a comparison of the compiled structures, not a controlled claim that every difference is caused only by depth." ] }, { "cell_type": "code", "execution_count": null, "id": "two-compile", "metadata": {}, "outputs": [], "source": [ "class TwoLayerRNN(brainstate.nn.Module):\n", " def __init__(self, n_in, n_rec, n_out):\n", " super().__init__()\n", " self.rnn1 = braintrace.nn.GRUCell(n_in, n_rec)\n", " self.rnn2 = braintrace.nn.GRUCell(n_rec, n_rec)\n", " self.out = braintrace.nn.Linear(n_rec, n_out)\n", "\n", " def update(self, x):\n", " h1 = self.rnn1(x)\n", " h2 = self.rnn2(h1)\n", " return self.out(h2)\n", "\n", "\n", "model2 = TwoLayerRNN(10, 32, 5)\n", "learner2 = braintrace.compile(model2, braintrace.D_RTRL, jnp.zeros(10))\n", "learner2.show_graph()" ] }, { "cell_type": "markdown", "id": "two-inspection-heading", "metadata": {}, "source": [ "### Inspect the Stacked Graph\n", "\n", "The compiler should keep the two sequential recurrent states in separate hidden groups. For each GRU layer, the update-gate (`Wz`) and candidate (`Wh`) weights have direct ETP relations to that layer's hidden group. Reset-gate (`Wr`) paths are excluded when they reach the hidden state through another trainable ETP primitive, while the readout remains non-temporal." ] }, { "cell_type": "code", "execution_count": null, "id": "two-inspection", "metadata": {}, "outputs": [], "source": [ "graph2 = learner2.graph\n", "\n", "print(f\"Number of hidden groups: {len(graph2.hidden_groups)}\")\n", "print(f\"Number of weight-hidden relations: {len(graph2.hidden_param_op_relations)}\")\n", "print(\"Hidden paths:\", [group.hidden_paths for group in graph2.hidden_groups])\n", "print(\"Report counts:\", learner2.report.counts)" ] }, { "cell_type": "markdown", "id": "comparison-heading", "metadata": {}, "source": [ "## Single-Layer vs Two-Layer\n", "\n", "The comparison below is generated from the two compiled learners rather than maintained as a separate hand-written result table. It distinguishes model architecture from the compiler structures discovered for each model." ] }, { "cell_type": "code", "execution_count": null, "id": "compiler-comparison", "metadata": {}, "outputs": [], "source": [ "single_counts = learner.report.counts\n", "two_counts = learner2.report.counts\n", "\n", "print(f\"{'Property':<32} {'Single layer':>14} {'Two layers':>14}\")\n", "print(\"-\" * 62)\n", "print(f\"{'Recurrent cell':<32} {'ValinaRNNCell':>14} {'GRUCell':>14}\")\n", "print(f\"{'Hidden groups':<32} {len(learner.graph.hidden_groups):>14} {len(learner2.graph.hidden_groups):>14}\")\n", "print(f\"{'Weight-hidden relations':<32} {len(learner.graph.hidden_param_op_relations):>14} {len(learner2.graph.hidden_param_op_relations):>14}\")\n", "print(f\"{'ETP weights':<32} {single_counts['etrace_weights']:>14} {two_counts['etrace_weights']:>14}\")\n", "print(f\"{'Excluded weights':<32} {single_counts['excluded_weights']:>14} {two_counts['excluded_weights']:>14}\")" ] }, { "cell_type": "markdown", "id": "comparison-interpretation", "metadata": {}, "source": [ "The second recurrent layer creates another independently represented hidden-state transition. The additional GRU gates also create more candidate parameter paths, so the relation-count difference must not be attributed to depth alone. In both models, the report explains classification outcomes and the graph identifies the exact parameter, primitive, and hidden-group connections behind them." ] }, { "cell_type": "markdown", "id": "combined-summary", "metadata": {}, "source": [ "## Summary\n", "\n", "The single-layer model establishes the inspection sequence; the two-layer model applies it without repeating the API explanation. Comparing the resulting reports and graphs shows why compiler decisions must be interpreted at the relation level rather than inferred from module names or parameter counts. These structural checks establish what the compiler selected, but they do not by themselves establish gradient correctness." ] } ], "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.11.8" } }, "nbformat": 4, "nbformat_minor": 5 }