{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "48c91def",
   "metadata": {},
   "source": [
    "# Training and evaluation\n",
    "\n",
    "Write an explicit optimization step, clip gradients, separate training from evaluation, and aggregate held-out metrics.\n",
    "\n",
    "These outputs were saved from local execution. The documentation renderer does not execute notebook cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4e1754c6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:46.376575Z",
     "iopub.status.busy": "2026-10-14T13:50:46.376478Z",
     "iopub.status.idle": "2026-10-14T13:50:46.975764Z",
     "shell.execute_reply": "2026-10-14T13:50:46.975202Z"
    },
    "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": "74755cad",
   "metadata": {},
   "source": [
    "## Prepare a controlled classification task\n",
    "\n",
    "The two-dimensional synthetic task makes the loop easy to inspect. Training and held-out examples use independent random draws. This is an API exercise, not a dataset benchmark."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5b1dd511",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:46.977281Z",
     "iopub.status.busy": "2026-10-14T13:50:46.977093Z",
     "iopub.status.idle": "2026-10-14T13:50:47.421983Z",
     "shell.execute_reply": "2026-10-14T13:50:47.421413Z"
    }
   },
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(8)\n",
    "x = rng.normal(size=(256,2)).astype(np.float32)\n",
    "y = (x[:,0] + .5*x[:,1] > 0).astype(np.int32)\n",
    "train = (jnp.asarray(x[:192]), jnp.asarray(y[:192]))\n",
    "valid = (jnp.asarray(x[192:]), jnp.asarray(y[192:]))\n",
    "class Classifier(nn.Module):\n",
    "    def __init__(self, *, rngs):\n",
    "        super().__init__()\n",
    "        self.hidden = nn.Linear(2,16,rngs=rngs)\n",
    "        self.dropout = nn.Dropout(.1)\n",
    "        self.head = nn.Linear(16,2,rngs=rngs)\n",
    "    def forward(self,x,*,rngs=None):\n",
    "        return self.head(self.dropout(F.relu(self.hidden(x)),rngs=rngs))\n",
    "model = Classifier(rngs=js.Rngs(params=8))\n",
    "optimizer = optim.Adam(model.parameters(),lr=.02)\n",
    "runtime = js.Rngs(dropout=9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75d39fa0",
   "metadata": {},
   "source": [
    "## Differentiate, clip, update\n",
    "\n",
    "The step returns gradients explicitly, then applies one optimizer update. Auxiliary predictions supply metrics without a second stochastic forward pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bf1a1128",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:47.423784Z",
     "iopub.status.busy": "2026-10-14T13:50:47.423624Z",
     "iopub.status.idle": "2026-10-14T13:50:48.469990Z",
     "shell.execute_reply": "2026-10-14T13:50:48.469402Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'correct': 191.0, 'count': 192.0, 'gradient_norm': 0.02445269748568535, 'loss': 0.02627979777753353}\n",
      "Optimizer updates: 100\n"
     ]
    }
   ],
   "source": [
    "@js.jit\n",
    "def train_step(model, optimizer, batch, rngs):\n",
    "    x,y = batch\n",
    "    def objective(candidate):\n",
    "        logits = candidate(x,rngs=rngs)\n",
    "        return F.cross_entropy(logits,y), (logits.argmax(-1)==y).sum()\n",
    "    (loss,correct), grads = js.value_and_grad(objective,has_aux=True)(model)\n",
    "    grads,norm = optim.clip_grad_norm(grads,1.)\n",
    "    optimizer.step(grads)\n",
    "    return {'loss':loss,'correct':correct,'count':y.size,'gradient_norm':norm}\n",
    "model.train()\n",
    "for step in range(100):\n",
    "    result = train_step(model,optimizer,train,runtime)\n",
    "print({name:float(value) for name,value in result.items()})\n",
    "print('Optimizer updates:',int(optimizer.step_count))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00d52bb3",
   "metadata": {},
   "source": [
    "## Evaluate without updating state\n",
    "\n",
    "Evaluation mode changes dropout and normalization behavior; it is not a gradient switch. Validation uses an explicit, separate function and a held-out batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "17b9e057",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:48.471453Z",
     "iopub.status.busy": "2026-10-14T13:50:48.471340Z",
     "iopub.status.idle": "2026-10-14T13:50:48.653266Z",
     "shell.execute_reply": "2026-10-14T13:50:48.652661Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Held-out loss: 0.031200522556900978\n",
      "Held-out accuracy: 1.0\n"
     ]
    }
   ],
   "source": [
    "@js.jit\n",
    "def evaluate(model,batch):\n",
    "    x,y = batch\n",
    "    logits = model(x)\n",
    "    return {'loss':F.cross_entropy(logits,y),\n",
    "            'correct':(logits.argmax(-1)==y).sum(),'count':y.size}\n",
    "model.eval()\n",
    "result = evaluate(model,valid)\n",
    "print('Held-out loss:',float(result['loss']))\n",
    "print('Held-out accuracy:',float(result['correct']/result['count']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fc2a6c6",
   "metadata": {},
   "source": [
    "## Aggregate by valid item count\n",
    "\n",
    "Do not average batch means with equal weight when the final batch is smaller. Masked sequence tasks should report valid-token counts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "94470e1d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:48.654411Z",
     "iopub.status.busy": "2026-10-14T13:50:48.654301Z",
     "iopub.status.idle": "2026-10-14T13:50:49.058049Z",
     "shell.execute_reply": "2026-10-14T13:50:49.057423Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.031184522405965254, 'accuracy': 1.0}\n"
     ]
    }
   ],
   "source": [
    "from jaxsaw.metrics import ClassificationAccumulator\n",
    "accumulator = ClassificationAccumulator()\n",
    "for start in range(0,len(valid[0]),15):\n",
    "    batch = tuple(array[start:start+15] for array in valid)\n",
    "    accumulator.update(evaluate(model,batch))\n",
    "print(accumulator.compute())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5ff85c5",
   "metadata": {},
   "source": [
    "## Adapt the objective\n",
    "\n",
    "Change the scalar objective for regression, reconstruction, or multiple supervised heads. In adversarial training, use separate optimizers and select the differentiated module explicitly. See the Fashion-MNIST example for both phases."
   ]
  }
 ],
 "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:49Z",
   "gpu_selection": "0",
   "kind": "api",
   "requires_figure": false,
   "source_sha256": "0e3671280366318c30f69a04b7a179649c5f5e254af15f9db3a45155fa3c7228",
   "source_rename": {
    "package_from": "neuraljaxlib",
    "package_to": "jaxsaw",
    "alias_from": "nj",
    "alias_to": "js"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}