{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9df7d2e1",
   "metadata": {},
   "source": [
    "# Modules and parameters\n",
    "\n",
    "Build reusable layers and residual blocks, inspect registered variables, and understand shared parameter ownership.\n",
    "\n",
    "These outputs were saved from local execution. The documentation renderer does not execute notebook cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "55e03538",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:38.962028Z",
     "iopub.status.busy": "2026-10-14T13:50:38.961934Z",
     "iopub.status.idle": "2026-10-14T13:50:39.543855Z",
     "shell.execute_reply": "2026-10-14T13:50:39.543156Z"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JAXsaw 0.1.0 · JAX 0.5.3\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CudaDevice(id=0)]\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import sys\n",
    "ROOT = Path.cwd() if (Path.cwd() / 'examples').exists() else Path.cwd().parent\n",
    "sys.path.insert(0, str(ROOT))\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import jaxsaw as js\n",
    "from jaxsaw import nn, optim, data, checkpoint, inspection\n",
    "from jaxsaw.nn import functional as F\n",
    "print(f\"JAXsaw {js.__version__} · JAX {jax.__version__}\")\n",
    "print(jax.devices())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "721298dc",
   "metadata": {},
   "source": [
    "## Define a residual module\n",
    "\n",
    "A module owns computation and state. The residual addition is ordinary array computation; it does not require a special model class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "640a2dfa",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:39.545122Z",
     "iopub.status.busy": "2026-10-14T13:50:39.544942Z",
     "iopub.status.idle": "2026-10-14T13:50:40.266349Z",
     "shell.execute_reply": "2026-10-14T13:50:40.265772Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output shape: (4, 2)\n",
      "Module | Type | Own parameters\n",
      "items.0 | Linear | 32\n",
      "items.1.hidden | Linear | 72\n",
      "items.1.output | Linear | 72\n",
      "items.1 | ResidualBlock | 0\n",
      "items.2 | Linear | 18\n",
      "<root> | Sequential | 0\n",
      "Total unique parameters: 194\n"
     ]
    }
   ],
   "source": [
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, width, *, rngs):\n",
    "        super().__init__()\n",
    "        self.hidden = nn.Linear(width, width, rngs=rngs)\n",
    "        self.output = nn.Linear(width, width, rngs=rngs)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self.output(F.relu(self.hidden(x)))\n",
    "\n",
    "model = nn.Sequential(nn.Linear(3, 8, rngs=js.Rngs(params=1)),\n",
    "                      ResidualBlock(8, rngs=js.Rngs(params=2)),\n",
    "                      nn.Linear(8, 2, rngs=js.Rngs(params=3)))\n",
    "x = jnp.ones((4, 3))\n",
    "print('Output shape:', model(x).shape)\n",
    "print(inspection.summary(model))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1962b4b3",
   "metadata": {},
   "source": [
    "## Inspect parameter paths\n",
    "\n",
    "Parameters are live variables. Use `.value` for their arrays. Stable paths allow inspection and strict state loading."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "44d52399",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:40.267800Z",
     "iopub.status.busy": "2026-10-14T13:50:40.267687Z",
     "iopub.status.idle": "2026-10-14T13:50:40.300468Z",
     "shell.execute_reply": "2026-10-14T13:50:40.299879Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "items.0.bias                        (8,)         float32\n",
      "items.0.weight                      (8, 3)       float32\n",
      "items.1.hidden.bias                 (8,)         float32\n",
      "items.1.hidden.weight               (8, 8)       float32\n",
      "items.1.output.bias                 (8,)         float32\n",
      "items.1.output.weight               (8, 8)       float32\n",
      "items.2.bias                        (2,)         float32\n",
      "items.2.weight                      (2, 8)       float32\n",
      "Snapshot arrays: 8\n"
     ]
    }
   ],
   "source": [
    "for name, parameter in model.named_parameters():\n",
    "    print(f'{name:35s} {str(parameter.value.shape):12s} {parameter.value.dtype}')\n",
    "snapshot = model.state_dict()\n",
    "print('Snapshot arrays:', len(snapshot))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab220cce",
   "metadata": {},
   "source": [
    "## Share weights deliberately\n",
    "\n",
    "Reusing the same module ties weights. The optimizer sees each parameter once, and gradients from every invocation accumulate into that shared variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3a08285c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:40.301692Z",
     "iopub.status.busy": "2026-10-14T13:50:40.301583Z",
     "iopub.status.idle": "2026-10-14T13:50:40.771074Z",
     "shell.execute_reply": "2026-10-14T13:50:40.770333Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unique parameter arrays: 2\n",
      "{'alias': 'model.items.0'}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shared-layer objective: 0.023256195709109306\n"
     ]
    }
   ],
   "source": [
    "shared = nn.Linear(8, 8, rngs=js.Rngs(params=4))\n",
    "tied = nn.Sequential(shared, nn.Tanh(), shared)\n",
    "print('Unique parameter arrays:', len(list(tied.parameters())))\n",
    "print(inspection.describe(tied)['root']['fields']['items'][2])\n",
    "loss, gradients = js.value_and_grad(lambda m: jnp.square(m(jnp.ones((2,8)))).mean())(tied)\n",
    "optimizer = optim.SGD(tied.parameters(), lr=.01)\n",
    "optimizer.step(gradients)\n",
    "print('Shared-layer objective:', float(loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50db9a54",
   "metadata": {},
   "source": [
    "## Return structured predictions\n",
    "\n",
    "Multiple outputs can be a dictionary or tuple. Keep the output structure fixed inside a compiled signature. The objective decides which outputs contribute to loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f1b88475",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:40.772716Z",
     "iopub.status.busy": "2026-10-14T13:50:40.772585Z",
     "iopub.status.idle": "2026-10-14T13:50:41.102633Z",
     "shell.execute_reply": "2026-10-14T13:50:41.101999Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'latent': (4, 4), 'logits': (4, 2)}\n"
     ]
    }
   ],
   "source": [
    "class EncoderWithHead(nn.Module):\n",
    "    def __init__(self, *, rngs):\n",
    "        super().__init__()\n",
    "        self.encoder = nn.Linear(3, 4, rngs=rngs)\n",
    "        self.head = nn.Linear(4, 2, rngs=rngs)\n",
    "    def forward(self, x):\n",
    "        latent = F.relu(self.encoder(x))\n",
    "        return {'latent': latent, 'logits': self.head(latent)}\n",
    "output = EncoderWithHead(rngs=js.Rngs(params=5))(x)\n",
    "print({name: array.shape for name,array in output.items()})"
   ]
  }
 ],
 "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.10.20"
  },
  "jaxsaw": {
   "executed_utc": "2026-10-14T13:50:41Z",
   "gpu_selection": "0",
   "kind": "api",
   "requires_figure": false,
   "source_sha256": "cf91d568c9086bda050d33f00e6629d4764378f09f15d2cceb447b26b61dd1cb",
   "source_rename": {
    "package_from": "neuraljaxlib",
    "package_to": "jaxsaw",
    "alias_from": "nj",
    "alias_to": "js"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}