{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "41fb09f6",
   "metadata": {},
   "source": [
    "# Checkpoints and randomness\n",
    "\n",
    "Save and restore model, optimizer, and random-stream state together. Verify the next update is reproduced, rather than checking only that weights load.\n",
    "\n",
    "These outputs were saved from local execution. The documentation renderer does not execute notebook cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f0739c77",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:21.864806Z",
     "iopub.status.busy": "2026-10-14T13:50:21.864716Z",
     "iopub.status.idle": "2026-10-14T13:50:22.515699Z",
     "shell.execute_reply": "2026-10-14T13:50:22.515134Z"
    },
    "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": "01719c2c",
   "metadata": {},
   "source": [
    "## Own random streams explicitly\n",
    "\n",
    "Initialization and runtime randomness use separate streams. The random stream is an explicit step argument and part of the checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2de7ee6b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:22.517213Z",
     "iopub.status.busy": "2026-10-14T13:50:22.517024Z",
     "iopub.status.idle": "2026-10-14T13:50:23.145361Z",
     "shell.execute_reply": "2026-10-14T13:50:23.144563Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial update: 6.502008438110352\n"
     ]
    }
   ],
   "source": [
    "model = nn.Linear(2,1,rngs=js.Rngs(params=15))\n",
    "optimizer = optim.Adam(model.parameters(),lr=.01)\n",
    "runtime = js.Rngs(sample=16)\n",
    "@js.jit\n",
    "def step(model,optimizer,rngs):\n",
    "    x = js.random.normal(rngs.next('sample'),(8,2))\n",
    "    target = x.sum(-1,keepdims=True)\n",
    "    loss,grads = js.value_and_grad(lambda m:F.mse_loss(m(x),target))(model)\n",
    "    optimizer.step(grads)\n",
    "    return loss\n",
    "print('Initial update:',float(step(model,optimizer,runtime)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55aa86a3",
   "metadata": {},
   "source": [
    "## Save a complete snapshot\n",
    "\n",
    "NPZ checkpoints contain arrays and a JSON manifest; loading does not execute embedded Python. Reconstruct the same architecture and optimizer configuration before restoration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b180df93",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:23.147109Z",
     "iopub.status.busy": "2026-10-14T13:50:23.146873Z",
     "iopub.status.idle": "2026-10-14T13:50:23.161427Z",
     "shell.execute_reply": "2026-10-14T13:50:23.160760Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Next update before restore: 2.25146222114563\n"
     ]
    }
   ],
   "source": [
    "directory = ROOT / 'runs' / 'api-checkpoint'\n",
    "directory.mkdir(parents=True,exist_ok=True)\n",
    "path = directory / 'state.npz'\n",
    "checkpoint.save(path,model=model,optimizer=optimizer,rngs=runtime,\n",
    "                progress={'next_update':int(optimizer.step_count)},metadata={'seed':15})\n",
    "expected_loss = step(model,optimizer,runtime)\n",
    "expected_weights = model.state_dict()\n",
    "print('Next update before restore:',float(expected_loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "304cc2fe",
   "metadata": {},
   "source": [
    "## Verify exact continuation on this backend\n",
    "\n",
    "Restoring RNG counters and optimizer moments is necessary to reproduce the next update. Array comparisons below verify actual continuation on this GPU/runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f9b9b19a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:23.162778Z",
     "iopub.status.busy": "2026-10-14T13:50:23.162566Z",
     "iopub.status.idle": "2026-10-14T13:50:23.175983Z",
     "shell.execute_reply": "2026-10-14T13:50:23.175394Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Next update after restore: 2.25146222114563\n",
      "Saved progress: {'next_update': 1}\n",
      "Continuation verified.\n"
     ]
    }
   ],
   "source": [
    "restored = checkpoint.restore(path,model=model,optimizer=optimizer,rngs=runtime)\n",
    "actual_loss = step(model,optimizer,runtime)\n",
    "np.testing.assert_allclose(actual_loss,expected_loss,atol=1e-7)\n",
    "for name,value in model.state_dict().items():\n",
    "    np.testing.assert_allclose(value,expected_weights[name],atol=1e-7)\n",
    "print('Next update after restore:',float(actual_loss))\n",
    "print('Saved progress:',restored.progress)\n",
    "print('Continuation verified.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2f211ec",
   "metadata": {},
   "source": [
    "## Weights-only snapshots\n",
    "\n",
    "Use state_dict when transferring weights into an already constructed matching model. It does not contain optimizer or RNG state, so it is insufficient for exact training continuation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fde08af3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:23.177345Z",
     "iopub.status.busy": "2026-10-14T13:50:23.177187Z",
     "iopub.status.idle": "2026-10-14T13:50:23.262813Z",
     "shell.execute_reply": "2026-10-14T13:50:23.262197Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Matching inference: [[-0.7157132]\n",
      " [-0.7157132]]\n"
     ]
    }
   ],
   "source": [
    "weights = model.state_dict()\n",
    "copy = nn.Linear(2,1,rngs=js.Rngs(params=99))\n",
    "copy.load_state_dict(weights)\n",
    "probe = jnp.ones((2,2))\n",
    "np.testing.assert_allclose(model(probe),copy(probe))\n",
    "print('Matching inference:',np.asarray(copy(probe)))"
   ]
  }
 ],
 "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:23Z",
   "gpu_selection": "0",
   "kind": "api",
   "requires_figure": false,
   "source_sha256": "b7b2e812852d4f4f7855670e9f096f02c4c423f5621f97aace7b5cfec716aa0b",
   "source_rename": {
    "package_from": "neuraljaxlib",
    "package_to": "jaxsaw",
    "alias_from": "nj",
    "alias_to": "js"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}