{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "085a3a23",
   "metadata": {},
   "source": [
    "# Data and batching\n",
    "\n",
    "Construct datasets, reproduce shuffles, handle incomplete batches, and resume an iterator without changing the sample order.\n",
    "\n",
    "These outputs were saved from local execution. The documentation renderer does not execute notebook cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "046c0279",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:24.300783Z",
     "iopub.status.busy": "2026-10-14T13:50:24.300693Z",
     "iopub.status.idle": "2026-10-14T13:50:24.979878Z",
     "shell.execute_reply": "2026-10-14T13:50:24.979060Z"
    },
    "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": "e674e1e7",
   "metadata": {},
   "source": [
    "## Construct a dataset\n",
    "\n",
    "A TensorDataset returns tuples with matching leading dimensions. DataLoader batches those samples into JAX arrays. Keep normalization statistics fitted on the training split."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8c3508c2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:24.981340Z",
     "iopub.status.busy": "2026-10-14T13:50:24.981150Z",
     "iopub.status.idle": "2026-10-14T13:50:24.985494Z",
     "shell.execute_reply": "2026-10-14T13:50:24.985031Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Examples: 10\n",
      "First sample: (array([0., 1., 2.], dtype=float32), np.int32(0))\n"
     ]
    }
   ],
   "source": [
    "features = np.arange(30,dtype=np.float32).reshape(10,3)\n",
    "targets = np.arange(10,dtype=np.int32)\n",
    "dataset = data.TensorDataset(features,targets)\n",
    "loader = data.DataLoader(dataset,batch_size=4,shuffle=True,seed=12)\n",
    "print('Examples:',len(dataset))\n",
    "print('First sample:',dataset[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "749cf98c",
   "metadata": {},
   "source": [
    "## Inspect an epoch\n",
    "\n",
    "The final batch is retained by default. Different shapes may compile separately; use drop_last for training or padding and a valid-item mask when appropriate. Never silently discard evaluation examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9cceea6f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:24.986655Z",
     "iopub.status.busy": "2026-10-14T13:50:24.986551Z",
     "iopub.status.idle": "2026-10-14T13:50:25.023829Z",
     "shell.execute_reply": "2026-10-14T13:50:25.023321Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch: (4, 3) sample IDs: [7 1 5 0]\n",
      "Batch: (4, 3) sample IDs: [3 4 2 6]\n",
      "Batch: (2, 3) sample IDs: [9 8]\n"
     ]
    }
   ],
   "source": [
    "ids = []\n",
    "for x,y in loader.iter_epoch(0):\n",
    "    print('Batch:',x.shape,'sample IDs:',np.asarray(y))\n",
    "    ids.extend(np.asarray(y).tolist())\n",
    "assert sorted(ids)==list(range(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5352e92e",
   "metadata": {},
   "source": [
    "## Reproduce shuffling\n",
    "\n",
    "Shuffle order is derived from the loader seed and epoch. Identical configuration gives identical order; a new epoch changes it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "95a58e63",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:25.024883Z",
     "iopub.status.busy": "2026-10-14T13:50:25.024732Z",
     "iopub.status.idle": "2026-10-14T13:50:25.029695Z",
     "shell.execute_reply": "2026-10-14T13:50:25.029284Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0: [7 1 5 0 3 4 2 6 9 8]\n",
      "Epoch 1: [1 0 8 3 7 2 5 4 6 9]\n"
     ]
    }
   ],
   "source": [
    "def order(epoch):\n",
    "    copy = data.DataLoader(dataset,batch_size=4,shuffle=True,seed=12)\n",
    "    return np.concatenate([np.asarray(y) for _,y in copy.iter_epoch(epoch)])\n",
    "print('Epoch 0:',order(0))\n",
    "print('Epoch 1:',order(1))\n",
    "assert np.array_equal(order(0),np.asarray(ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3e496c9",
   "metadata": {},
   "source": [
    "## Resume after a completed update\n",
    "\n",
    "The loader cursor advances before yielding. Save it only after the corresponding optimizer update completes. A restored iterator continues with the next batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2265ad0e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:25.030705Z",
     "iopub.status.busy": "2026-10-14T13:50:25.030609Z",
     "iopub.status.idle": "2026-10-14T13:50:25.034141Z",
     "shell.execute_reply": "2026-10-14T13:50:25.033684Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Restored next IDs: [3 4 2 6]\n"
     ]
    }
   ],
   "source": [
    "loader = data.DataLoader(dataset,batch_size=4,shuffle=True,seed=12)\n",
    "iterator = loader.iter_epoch(0)\n",
    "first = next(iterator)\n",
    "state = loader.state_dict()\n",
    "expected = next(iterator)[1]\n",
    "restored = data.DataLoader(dataset,batch_size=4,shuffle=True,seed=12)\n",
    "restored.load_state_dict(state)\n",
    "actual = next(restored.iter_epoch(0))[1]\n",
    "print('Restored next IDs:',np.asarray(actual))\n",
    "assert np.array_equal(expected,actual)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8df6daf",
   "metadata": {},
   "source": [
    "## Pad variable-length sequences\n",
    "\n",
    "Return explicit lengths and masks from a custom collator. Recurrent layers use lengths to freeze the carry; token losses use a valid-item mask to normalize correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f40cc179",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-14T13:50:25.035281Z",
     "iopub.status.busy": "2026-10-14T13:50:25.035189Z",
     "iopub.status.idle": "2026-10-14T13:50:25.724294Z",
     "shell.execute_reply": "2026-10-14T13:50:25.723648Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Padded shape: (2, 3, 1) lengths: [3 1]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequence: (2, 3, 4) hidden/cell: [(1, 2, 4), (1, 2, 4)]\n"
     ]
    }
   ],
   "source": [
    "def pad_sequences(samples):\n",
    "    sequences, labels = zip(*samples)\n",
    "    lengths = np.array([len(s) for s in sequences],dtype=np.int32)\n",
    "    padded = np.zeros((len(samples),int(lengths.max()),1),dtype=np.float32)\n",
    "    for i,sequence in enumerate(sequences):\n",
    "        padded[i,:len(sequence),0]=sequence\n",
    "    return jnp.asarray(padded),jnp.asarray(lengths),jnp.asarray(labels)\n",
    "batch = pad_sequences([(np.array([1.,2.,3.]),0),(np.array([4.]),1)])\n",
    "print('Padded shape:',batch[0].shape,'lengths:',batch[1])\n",
    "recurrent = nn.LSTM(1,4,rngs=js.Rngs(params=0))\n",
    "sequence,carry = recurrent(batch[0],lengths=batch[1])\n",
    "print('Sequence:',sequence.shape,'hidden/cell:',[a.shape for a in carry])"
   ]
  }
 ],
 "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:26Z",
   "gpu_selection": "0",
   "kind": "api",
   "requires_figure": false,
   "source_sha256": "eaabd4cb8003e6bf3d19ce50bbf052c26321bed6b220d75aeab7223d2fcccd24",
   "source_rename": {
    "package_from": "neuraljaxlib",
    "package_to": "jaxsaw",
    "alias_from": "nj",
    "alias_to": "js"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}