{
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "name": "aner",
      "display_name": "Aner",
      "language": "aner"
    },
    "language_info": {
      "name": "aner",
      "file_extension": ".aner"
    },
    "aner": {
      "execution_mode": "persistent-session",
      "runtime_mode": "development"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "# MNIST: functional neural training\n\nThe same MLP, immutable partitions, Xavier initialization, Adam optimizer and mini batch schedule are used at every level of control. Download the [verified preparation script](https://aner.albarena.ai/reference/tools/benchmarks/fetch_mnist.sh), then run `bash fetch_mnist.sh mnist` from this notebook's directory before execution. The data cells use that local `mnist` folder. The official test set is evaluated only after the fixed fifth epoch.\n\nSelect Dev in the Aner editor to see variables, training progress and gradient samples. Select Prod and restart/rerun every cell to measure execution without observations. Native Float64 is supported on CPU and CUDA; MPS requires Float32 or an explicit precision policy. Set the profile in the configuration cell before creating parameters.\n\nFor a longer cell set `aner.notebook.timeoutMs` to `600000` before execution. Changes preserve variables and affect subsequent cells; an actual timeout clears the session.\n"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Execution profile\n\nUse `cuda` on an available accelerator, or `mps` with `Float32` on a supported Mac. Placement is explicit and inspectable.\n"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "import aner.nn\n",
        "nn.configure(device: \"cpu\", precision: \"Float64\", policy: \"preserve\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Dataset protocol\n\n54,000 training rows, 6,000 stratified validation rows, and 10,000 untouched official test rows. No writer IDs are supplied by IDX; the validation split is row stratified and does not claim writer grouping.\n"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "// MNIST MLP: 784 -> 300 -> 100 -> 10. Fixed five epochs; test used once at the end.\n",
        "// First run tools/benchmarks/fetch_mnist.sh from the repository root.\n",
        "import aner.dataset\n",
        "import aner.ml\n",
        "import aner.dev\n",
        "import aner.viz\n",
        "import aner.metrics\n",
        "\n",
        "let seed = 42\n",
        "let epochs = 5\n",
        "let batch_size = 128\n",
        "let official = dataset.mnist(directory: \"mnist\", partition: \"train\")\n",
        "let testing = dataset.mnist(directory: \"mnist\", partition: \"test\")\n",
        "let partitions = official.partition(validation: 0.1, test: 0.0, seed: 2026, stratify: true, group_aware: false)\n",
        "let training = partitions.train()\n",
        "let validation = partitions.validation()\n",
        "print(partitions.audit())\n",
        "dataset.require_disjoint(training, testing)\n",
        "dataset.require_disjoint(validation, testing)\n",
        "training.require_training()\n",
        "// Placement follows CLI options; preserve is the default precision policy.\n",
        "print(nn.execution())\n",
        "print(nn.backend_details())\n",
        "dev.automatic(false)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "## Model and fixed training run\n\nKeep this phase in a single cell when measuring its elapsed time. The loss is the row weighted mean of each mini batch objective before its update; validation uses detached parameters after each epoch.\n"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "nn.reset_memory_peaks()\n",
        "let started = nn.clock()\n",
        "// Functional control: each Dense initialization and activation is visible.\n",
        "fn build_network(seed: Int64) {\n",
        "    return nn.sequential(seed: seed)\n",
        "        .dense(inputs: 784, outputs: 300, initialization: \"xavier\")\n",
        "        .relu()\n",
        "        .dense(inputs: 300, outputs: 100, initialization: \"xavier\")\n",
        "        .relu()\n",
        "        .dense(inputs: 100, outputs: 10, initialization: \"xavier\")\n",
        "}\n",
        "fn train_network(network: NeuralNetwork, samples: Dataset, held_out: Dataset,\n",
        "                 epochs: Int64, batch: Int64, seed: Int64) {\n",
        "    let contract = ml.training(epochs: epochs, batch: batch, rate: 0.001, seed: seed)\n",
        "        .validation(held_out).select_best(false)\n",
        "    let plan = ml.classifier(network, standardize: false, loss: \"cross_entropy\", optimizer: \"adam\")\n",
        "    return plan.fit(samples, contract, progress: true)\n",
        "}\n",
        "let fitted = train_network(build_network(seed), training, validation, epochs, batch_size, seed)\n",
        "\n",
        "print(\"BENCHMARK_FIT_SECONDS\")\n",
        "print(nn.clock() - started)\n",
        "print(\"BENCHMARK_TRAIN_LOSS\")\n",
        "print(fitted.training_loss())\n",
        "print(\"BENCHMARK_TRAIN_ACCURACY\")\n",
        "print(fitted.training_accuracy())\n",
        "let test_started = nn.clock()\n",
        "let evaluation = fitted.evaluate(testing)\n",
        "print(\"BENCHMARK_TEST_SECONDS\")\n",
        "print(nn.clock() - test_started)\n",
        "print(\"BENCHMARK_TEST_LOSS\")\n",
        "print(evaluation.loss())\n",
        "print(\"BENCHMARK_TEST_ACCURACY\")\n",
        "print(evaluation.accuracy())\n",
        "print(\"BENCHMARK_MEMORY\")\n",
        "print(nn.memory())\n",
        "print(fitted.summary())\n",
        "dev.capture(\"Final model\", fitted)\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}