{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "11763f53",
   "metadata": {},
   "source": [
    "# Inspection and custom layers\n",
    "\n",
    "Implement a differentiable custom layer, register non-trainable state, inspect its hierarchy, and use the pure state boundary.\n",
    "\n",
    "These outputs were saved from local execution. The documentation renderer does not execute notebook cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a28a9668",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:28.627479Z",
     "iopub.status.busy": "2026-10-14T13:50:28.627379Z",
     "iopub.status.idle": "2026-10-14T13:50:29.275062Z",
     "shell.execute_reply": "2026-10-14T13:50:29.274406Z"
    },
    "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": "3a5ea21b",
   "metadata": {},
   "source": [
    "## Implement an affine layer\n",
    "\n",
    "The layer uses ordinary JAX operations. Parameters receive gradients automatically; no manual backward method is needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "92580519",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:29.276491Z",
     "iopub.status.busy": "2026-10-14T13:50:29.276303Z",
     "iopub.status.idle": "2026-10-14T13:50:29.574769Z",
     "shell.execute_reply": "2026-10-14T13:50:29.574120Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Module | Type | Own parameters\n",
      "items.0 | Affine | 6\n",
      "items.1 | Linear | 4\n",
      "<root> | Sequential | 0\n",
      "Total unique parameters: 10\n"
     ]
    }
   ],
   "source": [
    "class Affine(nn.Module):\n",
    "    def __init__(self,features,*,rngs):\n",
    "        super().__init__()\n",
    "        self.features = features\n",
    "        self.scale = nn.Parameter(jnp.ones(features))\n",
    "        self.offset = nn.Parameter(jnp.zeros(features))\n",
    "    def forward(self,x):\n",
    "        return x*self.scale.value + self.offset.value\n",
    "model = nn.Sequential(Affine(3,rngs=js.Rngs(params=0)),\n",
    "                      nn.Linear(3,1,rngs=js.Rngs(params=1)))\n",
    "print(inspection.summary(model))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "446dafe2",
   "metadata": {},
   "source": [
    "## Inspect without mutating the live model\n",
    "\n",
    "output_spec runs an evaluation clone. The architecture description records ownership, shapes, configuration, and aliases; it is not a full operation-level execution graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9207db83",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:29.576163Z",
     "iopub.status.busy": "2026-10-14T13:50:29.576051Z",
     "iopub.status.idle": "2026-10-14T13:50:29.698055Z",
     "shell.execute_reply": "2026-10-14T13:50:29.697466Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'shape': [4, 1], 'dtype': 'float32'}\n",
      "{\n",
      "  \"schema_version\": 1,\n",
      "  \"scope\": \"module hierarchy; not an execution graph\",\n",
      "  \"root\": {\n",
      "    \"type\": \"jaxsaw.nn.module.Sequential\",\n",
      "    \"fields\": {\n",
      "      \"items\": [\n",
      "        {\n",
      "          \"type\": \"__main__.Affine\",\n",
      "          \"fields\": {\n",
      "            \"features\": 3,\n",
      "            \"scale\": {\n",
      "              \"kind\": \"Param\",\n",
      "              \"shape\": [\n",
      "                3\n",
      "              ],\n",
      "              \"dtype\": \"float32\"\n",
      "            },\n",
      "            \"offset\": {\n",
      "              \"kind\": \"Param\",\n",
      "              \"shape\": [\n",
      "                3\n",
      "              ],\n",
      "              \"dtype\": \"float32\"\n",
      "            }\n",
      "          }\n",
      "        },\n",
      "        {\n",
      "          \"type\": \"jaxsaw.nn.layers.Linear\",\n",
      "          \"fields\": {\n",
      "            \"in_features\": 3,\n",
      "            \"out_features\": 1,\n",
      "            \"weight\": {\n",
      "              \"kind\": \"Param\",\n",
      "              \"shape\": [\n",
      "                1,\n",
      "                3\n",
      "              ],\n",
      "              \"dtype\": \"float32\"\n",
      "            },\n",
      "            \"bias\": {\n",
      "              \"kind\": \"Param\",\n",
      "              \"shape\": [\n",
      "                1\n",
      "              ],\n",
      "              \"dtype\": \"float32\"\n",
      "            }\n",
      "          }\n",
      "        }\n",
      "      ]\n",
      "    }\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(inspection.output_spec(model,jnp.ones((4,3))))\n",
    "import json\n",
    "print(json.dumps(inspection.describe(model),indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4889393c",
   "metadata": {},
   "source": [
    "## Differentiate custom parameters\n",
    "\n",
    "Only Parameter leaves are differentiated by default. Buffers such as running statistics are persistent state but not optimizer parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d8eff0b7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:29.699134Z",
     "iopub.status.busy": "2026-10-14T13:50:29.699024Z",
     "iopub.status.idle": "2026-10-14T13:50:30.203575Z",
     "shell.execute_reply": "2026-10-14T13:50:30.202799Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loss: 0.08691398799419403\n",
      "Learned affine scale: [0.9986205 1.0018053 0.9978359]\n"
     ]
    }
   ],
   "source": [
    "loss,grads = js.value_and_grad(lambda m:jnp.square(m(jnp.ones((4,3)))).mean())(model)\n",
    "optimizer = optim.SGD(model.parameters(),lr=.01)\n",
    "optimizer.step(grads)\n",
    "print('Loss:',float(loss))\n",
    "print('Learned affine scale:',np.asarray(model[0].scale.value))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "215a3816",
   "metadata": {},
   "source": [
    "## Cross the pure functional boundary\n",
    "\n",
    "split separates static structure from dynamic state. apply returns output and updated state without mutating the original model. This is useful when integrating pure JAX computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "53f7ab9f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:30.205272Z",
     "iopub.status.busy": "2026-10-14T13:50:30.205097Z",
     "iopub.status.idle": "2026-10-14T13:50:30.324945Z",
     "shell.execute_reply": "2026-10-14T13:50:30.324349Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pure output: [[0.2679077]\n",
      " [0.2679077]]\n"
     ]
    }
   ],
   "source": [
    "graph,state = js.split(model)\n",
    "x = jnp.ones((2,3))\n",
    "output,next_state = js.apply(graph,state,x)\n",
    "np.testing.assert_allclose(output,model(x))\n",
    "print('Pure output:',np.asarray(output))\n",
    "js.update(model,next_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7c9e71d",
   "metadata": {},
   "source": [
    "## Register persistent buffers\n",
    "\n",
    "Use buffers for state you intend to checkpoint but not optimize. Side effects outside registered state, such as appending to a Python log, do not belong inside compiled steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5c07802b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:30.326523Z",
     "iopub.status.busy": "2026-10-14T13:50:30.326410Z",
     "iopub.status.idle": "2026-10-14T13:50:30.366738Z",
     "shell.execute_reply": "2026-10-14T13:50:30.366141Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Buffers: [('mean', (3,))]\n",
      "Trainable arrays: 0\n",
      "[[1. 2. 3.]]\n"
     ]
    }
   ],
   "source": [
    "class Center(nn.Module):\n",
    "    def __init__(self,mean):\n",
    "        super().__init__()\n",
    "        self.mean = nn.Buffer(jnp.asarray(mean,dtype=jnp.float32))\n",
    "    def forward(self,x):\n",
    "        return x-self.mean.value\n",
    "center = Center([1.,2.,3.])\n",
    "print('Buffers:',[(name,value.value.shape) for name,value in center.named_buffers()])\n",
    "print('Trainable arrays:',len(list(center.parameters())))\n",
    "print(center(jnp.array([[2.,4.,6.]])))"
   ]
  }
 ],
 "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:30Z",
   "gpu_selection": "0",
   "kind": "api",
   "requires_figure": false,
   "source_sha256": "b6a5de9e6bd4020c92f6b177097970af9ba2dcca580ce381a4c59042fb1d5560",
   "source_rename": {
    "package_from": "neuraljaxlib",
    "package_to": "jaxsaw",
    "alias_from": "nj",
    "alias_to": "js"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}