diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index cc5791a..d790a4a 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -1,14 +1,11 @@
# Contributing
-Thanks for taking a look. FastNN is meant to be read and edited, so every file
-is one idea and none of them are long. Most of what follows is about keeping it
-that way.
+FastNN is meant to be read and edited: every file is one idea and none are long.
+Most of this is about keeping it that way.
-If you want to work on something bigger than a bug fix, open an issue first so
-we can agree on the shape before you write it. Issues labelled
-`good first issue` are a decent place to start, and `help wanted` is for
-bigger work, mostly performance. Comment on an issue before you start on it
-so two people don't write the same fix.
+For anything bigger than a bug fix, open an issue first. `good first issue` is a
+good place to start, `help wanted` is for bigger work, mostly performance.
+Comment on an issue before you start so two people don't write the same fix.
During Hacktoberfest the [pinned issue](https://github.com/CanReader/FastNN/issues/56)
explains what gets the `hacktoberfest-accepted` label.
@@ -36,21 +33,18 @@ at that same toolkit. The [CUDA section of the README](README.md#cuda) has the
details, including `FASTNN_NVCC_CCBIN` for when `nvcc` rejects your system
compiler.
-Use `--release` for anything that trains, debug is around 50x slower.
+Use `--release` for anything that trains. Debug is around 50x slower.
## Pull requests
-1. Fork, then branch off `master`. Name the branch after what it does,
- `fix/batchnorm-variance` or `feat/conv3d`.
-2. Keep one change per PR. A bug fix and a refactor are two PRs.
-3. Fill in the template. "How it was verified" matters more than anything else
- in there, so be specific.
-4. CI has to be green before merge. PRs are squash merged, so the PR title
- becomes the commit on `master`. Keep it short, lowercase, and imperative:
+1. Fork and branch off `master`, e.g. `fix/batchnorm-variance`.
+2. One change per PR. A bug fix and a refactor are two PRs.
+3. Fill in the template. Be specific about how you verified it.
+4. CI has to be green. PRs are squash merged, so the title becomes the commit
+ on `master`: short, lowercase, imperative, like
`fix softmax for widths that are not a power of two`.
-`master` is protected. Nothing lands on it without a PR and a passing CI run,
-including my own changes.
+`master` is protected. Everything goes through a PR, including mine.
## What CI checks
@@ -64,13 +58,12 @@ including my own changes.
| Package | `cargo package`, so the published crate builds |
| CUDA build | compiles the kernels with nvcc and runs clippy with `--features cuda` |
-Hosted runners have no GPU. The CUDA job proves the kernels compile and link,
-but nothing in CI runs them. If you touch `cuda/`, run `cuda_parity` locally
-and paste the result in the PR. If you have no GPU, say so and I'll run it.
+Hosted runners have no GPU, so CI only checks that the kernels compile and
+link. If you touch `cuda/`, run `cuda_parity` locally and paste the result in
+the PR. No GPU? Say so and I'll run it.
A nightly job reruns the suite in release on stable and beta and trains the
-small examples, so problems that only show up with optimizations still get
-caught.
+small examples.
## Code conventions
@@ -105,12 +98,12 @@ for an activation is in `tensor/ops/activation.rs` and its derivative is in
5. If it has a CUDA kernel, declare it in `cuda/ffi.rs`, wrap it in
`cuda/kernels.rs`, and add a case to `tests/cuda_parity.rs`.
-The README has a full worked example under "Extending it".
+The README has a worked example under "Extending".
## Writing a gradcheck case
Gradcheck compares each backward rule against central finite differences of its
-own forward. Two things trip people up:
+own forward. Two things to watch:
* **Keep values bounded.** Finite differences subtract two nearby losses. If the
loss is in the thousands, `f32` loses most of its digits to cancellation and
@@ -141,15 +134,13 @@ git push origin v0.4.0
```
The release workflow checks the tag against `Cargo.toml`, runs the suite in
-release, builds the example programs for Linux (x86_64, ARM64, and x86_64 with
-CUDA), macOS (Intel and Apple Silicon), and Windows, publishes to crates.io, and
-creates the GitHub release with the archives, a `SHA256SUMS` file, and notes
-from the merged PRs. Release notes are grouped by PR label, so label your PRs.
+release, builds the examples for Linux (x86_64, ARM64, x86_64 with CUDA), macOS
+(Intel and Apple Silicon) and Windows, publishes to crates.io, and creates the
+GitHub release with the archives, `SHA256SUMS` and notes from the merged PRs.
+Notes are grouped by PR label, so label your PRs.
-Any PR that touches `release.yml` runs the whole thing except the publish
-steps, so the pipeline gets tested before a real tag depends on it.
+A PR that touches `release.yml` runs everything except the publish steps.
## License
-By contributing you agree that your work is released under the MIT license,
-same as the rest of the project.
+Contributions are released under the MIT license, same as the project.
diff --git a/README.md b/README.md
index 3261f08..d6a9d50 100644
--- a/README.md
+++ b/README.md
@@ -1,8 +1,6 @@
-
-
# FastNN
-**A deep learning library in Rust, with CUDA kernels for the parts that matter.**
+A deep learning library in Rust, with CUDA kernels for the parts that matter.
[](https://github.com/CanReader/FastNN/actions/workflows/ci.yml)
[](https://crates.io/crates/fastnn)
@@ -10,64 +8,30 @@
[](https://github.com/CanReader/FastNN/blob/master/Cargo.toml)
[](LICENSE)
-[Quick start](#quick-start) ·
-[Examples](#examples) ·
-[API docs](https://docs.rs/fastnn) ·
-[Releases](https://github.com/CanReader/FastNN/releases) ·
-[Contributing](CONTRIBUTING.md)
-
-**Hacktoberfest 2026:** FastNN is taking part. [Start here](https://github.com/CanReader/FastNN/issues/56).
-
-
-
-FastNN is a from-scratch deep learning library: tensors with reverse-mode
-autograd, the usual layers, transformers with KV-cached generation, modern
-optimizers, data loading, resumable checkpoints, and safetensors interchange.
-It runs on the CPU out of the box and on NVIDIA GPUs with one feature flag.
-
-It is built to be read. Every file is one idea and none of them are long, so you
-can follow a `loss.backward()` call from the training loop down to the kernel
-and back without getting lost.
-
-* **Readable.** 86 source files, about 12k lines. Forward ops and their backward
- rules live in mirrored files, one per category.
-* **Verified.** Every backward rule is checked against finite differences of its
- own forward, and every CUDA kernel against the CPU path.
-* **Built for real runs.** Exact resume including optimizer state, NaN tracing
- to the op that produced it, frozen layers that skip their gradient entirely.
-
-## Contents
-
-* [Features](#features)
-* [Installation](#installation)
-* [Quick start](#quick-start)
-* [Examples](#examples)
-* [Platform support](#platform-support)
-* [How it fits together](#how-it-fits-together)
-* [Long training runs](#long-training-runs)
-* [Extending FastNN](#extending-fastnn)
-* [Error handling](#error-handling)
-* [Testing](#testing)
-* [CUDA](#cuda)
-* [Project layout](#project-layout)
-* [Limitations](#limitations)
-* [Contributing](#contributing)
-* [License](#license)
+Tensors with reverse-mode autograd, the usual layers, transformers with
+KV-cached generation, optimizers, data loading, resumable checkpoints and
+safetensors. It runs on the CPU by default and on NVIDIA GPUs with the `cuda`
+feature.
+
+The codebase is small on purpose: 86 source files, about 12k lines, one idea
+per file. You can follow `loss.backward()` from the training loop down to the
+kernel. Every backward rule is checked against finite differences of its own
+forward, and every CUDA kernel against the CPU path.
## Features
| Area | What is in it |
|---|---|
-| **Tensors** | Broadcasting arithmetic, matmul (batched, transposed forms), reductions, indexing, views, `cat`/`stack`/`narrow`, CPU or CUDA storage |
-| **Autograd** | Reverse mode with no tape and no global state, `no_grad`, `detect_anomaly`, gradient clipping by norm or value |
-| **Layers** | `Linear`, `Conv1d`, `Conv2d` (dilation, groups, depthwise), `ConvTranspose1d/2d`, pooling, `BatchNorm2d`, `LayerNorm`, `RMSNorm`, `Dropout`, `Embedding`, `LSTM`, `GRU` |
-| **Transformers** | `MultiHeadAttention` with causal and padding masks, encoder and decoder stacks, full encoder-decoder `Transformer` |
-| **Generation** | `KvCache` for incremental decoding, `Sampler` (temperature, top-k, top-p, repetition penalty), `BeamSearch` |
-| **Optimizers** | `SGD`, `Adam`, `AdamW`, `RMSprop`, `Adagrad`, `Adadelta`, `RAdam`, `Lion`, plus `Lookahead` and `Ema` wrappers |
-| **Schedules** | `Warmup`, `CosineAnnealing`, `OneCycle`, `StepDecay`, `Constant` |
-| **Losses** | Cross-entropy (label smoothing, class weights, ignore index), MSE, MAE, BCE, Huber, KL, NLL, focal, dice, Gaussian and Poisson NLL, triplet, contrastive, InfoNCE |
-| **Data** | `Dataset` trait, shuffling `DataLoader` with device placement, built-in MNIST |
-| **Checkpoints** | Native weights, full training state (weights + optimizer + step), safetensors read and write with F16/BF16 widening |
+| Tensors | Broadcasting arithmetic, matmul (batched, transposed forms), reductions, indexing, views, `cat`/`stack`/`narrow`, CPU or CUDA storage |
+| Autograd | Reverse mode, no tape, no global state. `no_grad`, `detect_anomaly`, gradient clipping by norm or value |
+| Layers | `Linear`, `Conv1d`, `Conv2d` (dilation, groups, depthwise), `ConvTranspose1d/2d`, pooling, `BatchNorm2d`, `LayerNorm`, `RMSNorm`, `Dropout`, `Embedding`, `LSTM`, `GRU` |
+| Transformers | `MultiHeadAttention` with causal and padding masks, encoder and decoder stacks, encoder-decoder `Transformer` |
+| Generation | `KvCache`, `Sampler` (temperature, top-k, top-p, repetition penalty), `BeamSearch` |
+| Optimizers | `SGD`, `Adam`, `AdamW`, `RMSprop`, `Adagrad`, `Adadelta`, `RAdam`, `Lion`, plus `Lookahead` and `Ema` wrappers |
+| Schedules | `Warmup`, `CosineAnnealing`, `OneCycle`, `StepDecay`, `Constant` |
+| Losses | Cross-entropy (label smoothing, class weights, ignore index), MSE, MAE, BCE, Huber, KL, NLL, focal, dice, Gaussian and Poisson NLL, triplet, contrastive, InfoNCE |
+| Data | `Dataset` trait, shuffling `DataLoader` with device placement, MNIST |
+| Checkpoints | Weights, full training state (weights, optimizer, step), safetensors read and write with F16/BF16 widening |
## Installation
@@ -76,14 +40,12 @@ cargo add fastnn # CPU only, no CUDA toolkit needed
cargo add fastnn --features cuda # with the CUDA kernels
```
-| Requirement | Version |
-|---|---|
-| Rust | 1.87 or newer |
-| CUDA toolkit (only with `--features cuda`) | 11.8 or newer, for GPUs from Turing to Hopper (compute capability 7.5 to 9.0) |
+Needs Rust 1.87 or newer. With `--features cuda` you also need the CUDA toolkit,
+11.8 or newer, for GPUs from Turing to Hopper (compute capability 7.5 to 9.0).
## Quick start
-A complete MNIST training loop:
+An MNIST training loop:
```rust
use fastnn::prelude::*;
@@ -111,8 +73,7 @@ for batch in loader.iter() {
save(&model, "mnist.fdl")?;
```
-No gradient-mode flags, no borrow dance around the optimizer. That is the whole
-training loop. To train on a GPU, move the model and the loader:
+To train on a GPU, move the model and the loader:
```rust
let device = Device::best(); // GPU if there is one, else CPU
@@ -127,21 +88,21 @@ cargo run --release --example simple_mlp
cargo run --release --example mnist_cnn --features cuda
```
-Always use `--release`. Debug builds are roughly 50x slower.
+Use `--release`. Debug builds are about 50x slower.
| Example | What it shows |
|---|---|
-| `simple_mlp` | The smallest complete training loop, an MLP learning XOR |
+| `simple_mlp` | MLP learning XOR, the smallest training loop |
| `mnist_mlp` | Dataset download, batching, evaluation, checkpointing |
| `mnist_cnn` | Convolutions, batch norm, pooling, an LR schedule |
-| `char_lm` | A small GPT trained from scratch that generates text, with resume on Ctrl-C |
+| `char_lm` | Small GPT trained from scratch, resumes after Ctrl-C |
| `finetune` | Pretrain, freeze the backbone, train a new head |
| `gan` | Two networks trained against each other |
-| `vae` | A variational autoencoder on 2-D data |
-| `reinforce` | Policy-gradient reinforcement learning on a bandit |
-| `multi_task` | One shared trunk, two heads, one combined loss |
+| `vae` | Variational autoencoder on 2-D data |
+| `reinforce` | Policy gradient on a bandit |
+| `multi_task` | Shared trunk, two heads, one combined loss |
-`char_lm` trains on a small embedded corpus by default, or on any text file:
+`char_lm` uses a small embedded corpus by default, or any text file:
```bash
curl -o shakespeare.txt \
@@ -155,9 +116,8 @@ and watch the network light up as it predicts.
### Prebuilt binaries
-Every [release](https://github.com/CanReader/FastNN/releases) ships all of the
-examples prebuilt, so you can try FastNN without installing Rust. Download the
-archive for your platform, extract it, and run any example directly.
+Each [release](https://github.com/CanReader/FastNN/releases) has the examples
+prebuilt, so you can try them without installing Rust.
| Platform | Archive |
|---|---|
@@ -168,10 +128,9 @@ archive for your platform, extract it, and run any example directly.
| macOS Intel | `fastnn--x86_64-macos.tar.gz` |
| Windows x86_64 | `fastnn--x86_64-windows.zip` |
-Each release includes a `SHA256SUMS` file and a build provenance attestation
-for every archive, which you can check with
-`gh attestation verify --repo CanReader/FastNN`. The CUDA build links
-against the CUDA 13 runtime, so it needs that installed.
+Releases include a `SHA256SUMS` file and a build provenance attestation for each
+archive (`gh attestation verify --repo CanReader/FastNN`). The CUDA
+build links against the CUDA 13 runtime.
## Platform support
@@ -179,7 +138,7 @@ against the CUDA 13 runtime, so it needs that installed.
|---|---|---|
| Linux x86_64 | Tested in CI | Built in CI, parity tested on real hardware |
| Linux ARM64 | Built for releases | Not supported |
-| macOS (Intel, Apple Silicon) | Tested in CI | Not available on macOS |
+| macOS (Intel, Apple Silicon) | Tested in CI | Not available |
| Windows x86_64 | Tested in CI | Should build, not tested |
## How it fits together
@@ -197,34 +156,25 @@ against the CUDA 13 runtime, so it needs that installed.
optim opt.step() reads those slots and updates in place
```
-Everything starts on the CPU. `model.to_device(device)` moves a model,
-`DataLoader::to_device` moves its batches, and mixing devices panics rather
-than copying behind your back.
+Everything starts on the CPU. `model.to_device(device)` and
+`DataLoader::to_device` move things over, and mixing devices panics instead of
+copying behind your back.
-Two design decisions shape the rest of the library.
+The graph is the tensors. A tensor produced by a differentiable op carries a
+node naming the rule that made it and the tensors it consumed, and
+`loss.backward()` walks that. Drop the loss and the graph frees itself. A tensor
+with no node is plain data, so there is no `requires_grad` flag. Wrap inference
+in `no_grad(|| ...)` to skip building a graph.
-**The graph is the tensors.** There is no tape and no global state. A tensor
-produced by a differentiable op carries a node naming the rule that made it and
-the tensors it consumed, so `loss.backward()` just walks that structure. Drop the
-loss and the graph frees itself. A tensor with no node is plain data, so there is
-no `requires_grad` flag to keep in sync. Wrap inference in `no_grad(|| ...)` to
-skip building a graph at all.
-
-**Parameters are shared slots.** `Param` is a handle, not a value. The model and
-the optimizer hold handles to the same weight and the same gradient, which is why
-`opt.step()` needs no borrow of the model. It is also why weight tying works
-without special support: use the same handle twice and both paths' gradients add
-up on their own.
+`Param` is a handle, not a value. The model and the optimizer hold handles to
+the same weight and gradient, so `opt.step()` doesn't borrow the model. Weight
+tying works the same way: use the handle twice and the gradients add up.
## Long training runs
-Four things a long, unattended run needs that the quick start does not show.
-
-**Resume exactly where you stopped.** `save()` writes weights, which is right for
-a finished model. Resuming also needs the optimizer: momentum, Adam's moment
-estimates, and the step count that drives bias correction. Restore weights alone
-and the optimizer restarts at step 1, so the first update after resuming lands
-far harder than it should.
+**Resume.** `save()` writes weights, which is enough for a finished model.
+Resuming also needs the optimizer's moments and step count. Without them Adam
+restarts at step 1 and the first update is far too large.
```rust
let mut step = load_training(&model, &mut opt, "run.fdl").unwrap_or(0);
@@ -239,9 +189,9 @@ while step < total {
Loading a plain checkpoint with `load_training` is an error, not a silent
optimizer reset.
-**Find the NaN at the op that made it.** One non-finite gradient becomes a
-non-finite weight, and every activation downstream is NaN from then on. By the
-time the loss prints NaN the checkpoint is poisoned and the culprit is long gone.
+**NaNs.** One non-finite gradient poisons the weights, and by the time the loss
+prints NaN the culprit is long gone. `detect_anomaly` panics at the op that
+produced it.
```rust
detect_anomaly(|| {
@@ -251,20 +201,18 @@ detect_anomaly(|| {
if !opt.gradients_are_finite() { continue; } // cheap guard: skip the batch
```
-Both read every gradient, so they are debugging and guard-rail tools, not
-something to leave on in a healthy loop.
+Both read every gradient, so don't leave them on in a healthy loop.
-**Freeze a backbone.** A frozen parameter hands out a detached tensor, so the
-backward pass stops there and no gradient is computed only to be thrown away.
+**Freezing.** A frozen parameter hands out a detached tensor, so the backward
+pass stops there and no gradient is computed for it.
```rust
backbone.freeze();
let mut opt = Adam::new(model.parameters(), 1e-3); // updates only the head
```
-**Report a bad shape instead of crashing.** Ordinary ops panic, which is right
-inside a model. At the edge of an application, where a shape comes from a config
-file or an upload, use the `try_` variants.
+**Fallible ops.** Ordinary ops panic on bad shapes. Where a shape comes from a
+config file or an upload, use the `try_` variants.
```rust
let y = x.try_matmul(&w)?; // Error::Shape, not a panic
@@ -272,13 +220,13 @@ x.try_reshape(&[2, -1])?;
table.try_index_select(&ids)?;
```
-They validate and then delegate, so there is still exactly one implementation of
-each operation, and they build the same graph.
+They validate and then call the panicking op, so there is one implementation of
+each operation.
-## Extending FastNN
+## Extending
-**A new layer.** Implement `forward`, plus `named_parameters` if it has weights.
-Device placement, parameter counting, and gradient zeroing all come from those.
+**A layer.** Implement `forward`, plus `named_parameters` if it has weights.
+Device placement, parameter counting and gradient zeroing follow from those.
```rust
struct Residual { inner: Linear }
@@ -293,9 +241,8 @@ impl Module for Residual {
}
```
-**A new op.** Write the forward in `tensor/ops/`, the backward rule in
-`autograd/ops/`, attach it with `with_grad`, and add a case to
-`tests/gradcheck.rs`.
+**An op.** Forward in `tensor/ops/`, backward rule in `autograd/ops/`, attach it
+with `with_grad`, add a case to `tests/gradcheck.rs`.
```rust
// tensor/ops/unary.rs
@@ -313,23 +260,18 @@ impl Backward for SoftplusBackward {
}
```
-The rule returns gradients in the same order as the inputs. Save detached
-tensors: a saved value that still carries its own node would pin the graph that
-produced it.
-
-**A new dataset.** Implement `Dataset` and hand it to a `DataLoader`. Items are
-written into caller-provided slices rather than returned as tensors, so batching
-60,000 images does not build 60,000 throwaway ones.
-
-[CONTRIBUTING.md](CONTRIBUTING.md) has the full checklist for adding an op.
+The rule returns gradients in input order. Save detached tensors, otherwise the
+saved value pins the graph that produced it. [CONTRIBUTING.md](CONTRIBUTING.md)
+has the full checklist.
-## Error handling
+**A dataset.** Implement `Dataset` and pass it to a `DataLoader`. Items are
+written into caller-provided slices instead of returned as tensors, so batching
+60,000 images doesn't allocate 60,000 throwaway ones.
-Tensor maths panics on shape and device mistakes. Those are bugs in the calling
-code, like indexing past the end of a slice, and threading `Result` through every
-`add` would bury the model in `?` for no safety gained.
+## Errors
-`Result` is for what genuinely fails at runtime:
+Shape and device mistakes panic. They are bugs in the calling code, like
+indexing past the end of a slice. `Result` is for things that fail at runtime:
```rust
let device = Device::cuda(0)?; // no GPU
@@ -350,42 +292,32 @@ cargo test --features cuda --test cuda_parity -- --test-threads=1 # CPU vs GPU
cargo bench # throughput
```
-`gradcheck` compares every backward rule with central finite differences of its
-own forward. It is why the autograd layer can be trusted, and a new rule without
-a case there is not finished.
+`cuda_parity` runs each op on both devices and compares. It skips itself when
+there is no GPU. It has already caught a real bug: the softmax kernel's block
+reduction assumed a power-of-two thread count and dropped the tail of every row
+whose width wasn't one, which included MNIST's ten classes.
-`cuda_parity` runs each op on both devices and compares the results. It skips
-itself when there is no GPU, so a CPU-only machine still gets a green run. It has
-already caught a real bug: the softmax kernel's block reduction assumed a
-power-of-two thread count and silently dropped the tail of every row whose width
-was not one, which included MNIST's ten classes.
-
-Every pull request runs formatting, clippy, the full suite on Linux, macOS, and
-Windows, an MSRV build, docs, packaging, and a CUDA compile. A nightly job reruns
-everything in release mode on stable and beta and trains the examples end to end.
+CI runs formatting, clippy, the suite on Linux, macOS and Windows, an MSRV
+build, docs, packaging and a CUDA compile. A nightly job reruns the suite in
+release on stable and beta and trains the examples that don't need a download.
## CUDA
-The `cuda` feature is opt-in, so a plain `cargo add fastnn` never needs the
-toolkit. With the feature on, `build.rs` compiles `cuda/kernels.cu` with `nvcc`
-and links `cudart`, `cublas`, and `curand`.
+The `cuda` feature is opt-in. With it on, `build.rs` compiles `cuda/kernels.cu`
+with `nvcc` and links `cudart`, `cublas` and `curand`.
-* `nvcc` is found on `PATH`, so put the toolkit's `bin` directory there, or
- point the `NVCC` environment variable at the compiler directly.
+* `nvcc` is found on `PATH`, or set the `NVCC` environment variable.
* Headers and libraries come from `CUDA_PATH`, then `CUDA_HOME`, then the
- default install location. Make sure that is the same toolkit as the `nvcc`
- you picked, or the kernels build against one CUDA version and link against
- another.
+ default install location. Make sure that is the same toolkit as your `nvcc`.
* If `nvcc` rejects your system compiler as too new, point `FASTNN_NVCC_CCBIN`
at one it accepts.
* Without the feature, `cuda/stubs.c` supplies the symbols, `Device::cuda(0)`
returns an error, and everything runs on the CPU.
-Two things carry most of the GPU performance. Matrix multiplication goes through
-cuBLAS, and the two transposed forms the backward pass needs (`matmul_nt`,
-`matmul_tn`) use a transpose flag instead of building a transposed copy. GPU
-allocations come from a size-keyed free list, so the many short-lived temporaries
-a training step creates cost a hash lookup instead of a driver round trip.
+Matmul goes through cuBLAS. The transposed forms the backward pass needs
+(`matmul_nt`, `matmul_tn`) use a transpose flag instead of a transposed copy.
+GPU allocations come from a size-keyed free list, so the short-lived temporaries
+in a training step don't each hit the driver.
## Project layout
@@ -424,35 +356,30 @@ src/
cuda/kernels.cu every GPU kernel
```
-`tensor/ops/` and `autograd/ops/` mirror each other file for file: the forward
-for `relu` is in `tensor/ops/activation.rs` and its derivative is in
+`tensor/ops/` and `autograd/ops/` mirror each other file for file. The forward
+for `relu` is in `tensor/ops/activation.rs`, its derivative in
`autograd/ops/activation.rs`.
## Limitations
-* Compute is `f32`. Checkpoints stored as F16 or BF16 load fine and are widened
- on read, but the maths runs in single precision.
-* `LSTM` and `GRU` are built from ordinary differentiable ops, one graph node per
- gate per timestep. Correct, and fine for short sequences; use a transformer for
- long ones.
-* Tensors are always contiguous. `permute` and `expand` write a new buffer rather
- than returning a strided view.
+* Compute is `f32`. F16 and BF16 checkpoints load and are widened on read.
+* `LSTM` and `GRU` are built from ordinary ops, one graph node per gate per
+ timestep. Fine for short sequences, use a transformer for long ones.
+* Tensors are always contiguous. `permute` and `expand` write a new buffer
+ instead of returning a strided view.
## Contributing
-Bug reports, ideas, and pull requests are all welcome. [CONTRIBUTING.md](CONTRIBUTING.md)
-covers the setup, the conventions, what CI checks, and how to add an op. Issues
-labelled [`good first issue`](https://github.com/CanReader/FastNN/labels/good%20first%20issue)
-are a good place to start. Each one points at the file and line and says what a
-fix looks like.
+See [CONTRIBUTING.md](CONTRIBUTING.md). Issues labelled
+[`good first issue`](https://github.com/CanReader/FastNN/labels/good%20first%20issue)
+are a good place to start. Report security problems privately, see
+[SECURITY.md](SECURITY.md).
During Hacktoberfest, read the [pinned issue](https://github.com/CanReader/FastNN/issues/56)
first. It explains how claiming an issue works and what counts.
If FastNN is useful to you, a star helps other people find it.
-Please report security problems privately, as described in [SECURITY.md](SECURITY.md).
-
## License
-FastNN is released under the [MIT License](LICENSE).
+[MIT](LICENSE)
diff --git a/benches/tensor_ops.rs b/benches/tensor_ops.rs
index 3115a29..f183a91 100644
--- a/benches/tensor_ops.rs
+++ b/benches/tensor_ops.rs
@@ -1,15 +1,9 @@
-//! Throughput benchmarks for the hot paths.
-//!
-//! cargo bench # CPU
-//! cargo bench --features cuda # CUDA, if the toolkit is present
-//!
-//! Benchmarks run on whichever device is available, so the same command measures
-//! either backend. GPU timings include the launch, not just the kernel.
+//! Hot path throughput on the GPU if there is one, otherwise the CPU.
+//! GPU timings include the launch.
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use fastnn::prelude::*;
-/// The device under test, named for the report.
fn device() -> (Device, &'static str) {
match Device::cuda(0) {
Ok(device) => (device, "cuda"),
@@ -25,7 +19,7 @@ fn matmul(c: &mut Criterion) {
let a = Tensor::randn(&[size, size]).to(device);
let b = Tensor::randn(&[size, size]).to(device);
- // Two flops per multiply-accumulate.
+ // Two flops per multiply-add.
group.throughput(Throughput::Elements(2 * (size * size * size) as u64));
group.bench_with_input(BenchmarkId::from_parameter(size), &size, |bench, _| {
bench.iter(|| a.matmul(&b))
@@ -34,9 +28,7 @@ fn matmul(c: &mut Criterion) {
group.finish();
}
-/// The transpose-fused forms, which the backward pass of every linear layer
-/// leans on. They should track plain matmul closely — if they do not, a
-/// transposed copy is being materialised somewhere it should not be.
+/// Should track plain matmul, if not a transposed copy is being materialised.
fn matmul_fused(c: &mut Criterion) {
let (device, name) = device();
let mut group = c.benchmark_group(format!("matmul_fused/{name}"));
@@ -82,7 +74,6 @@ fn reductions(c: &mut Criterion) {
group.finish();
}
-/// Rearranging is pure memory traffic, and it sits in every attention forward.
fn layout(c: &mut Criterion) {
let (device, name) = device();
let mut group = c.benchmark_group(format!("layout/{name}"));
@@ -97,7 +88,6 @@ fn layout(c: &mut Criterion) {
group.finish();
}
-/// A whole training step, which is what actually matters.
fn training_step(c: &mut Criterion) {
let (device, name) = device();
let mut group = c.benchmark_group(format!("step/{name}"));
diff --git a/build.rs b/build.rs
index 29d903c..383f325 100644
--- a/build.rs
+++ b/build.rs
@@ -18,7 +18,6 @@ fn build_cpu_stubs() {
}
fn build_cuda() {
- // Find CUDA toolkit
let cuda_path = env::var("CUDA_PATH")
.or_else(|_| env::var("CUDA_HOME"))
.unwrap_or_else(|_| {
@@ -36,7 +35,6 @@ fn build_cuda() {
PathBuf::from(&cuda_path).join("lib64")
};
- // Compile CUDA kernels using cc with nvcc
let mut build = cc::Build::new();
build
.cuda(true)
@@ -53,9 +51,7 @@ fn build_cuda() {
.include(&cuda_include)
.file("cuda/kernels.cu");
- // nvcc rejects host compilers newer than it knows. FASTNN_NVCC_CCBIN picks
- // one explicitly; otherwise fall back to g++-15 when the system g++ is too
- // new and g++-15 is around.
+ // nvcc rejects host compilers newer than it knows. Take FASTNN_NVCC_CCBIN, else g++-15 if installed.
println!("cargo:rerun-if-env-changed=FASTNN_NVCC_CCBIN");
if let Ok(ccbin) = env::var("FASTNN_NVCC_CCBIN") {
build.flag(format!("-ccbin={ccbin}"));
@@ -69,7 +65,6 @@ fn build_cuda() {
build.compile("fastnn_cuda_kernels");
- // Link CUDA runtime and libraries
println!("cargo:rustc-link-search=native={}", cuda_lib.display());
println!("cargo:rustc-link-lib=dylib=cudart");
println!("cargo:rustc-link-lib=dylib=cublas");
diff --git a/cuda/include/kernels.h b/cuda/include/kernels.h
index eb0b150..28e14a3 100644
--- a/cuda/include/kernels.h
+++ b/cuda/include/kernels.h
@@ -7,17 +7,11 @@
extern "C" {
#endif
-// ============================================================================
-// Device Management
-// ============================================================================
int fastnn_cuda_init(int device_id);
int fastnn_cuda_device_count();
int fastnn_cuda_synchronize();
void fastnn_cuda_get_memory_info(size_t* free_mem, size_t* total_mem);
-// ============================================================================
-// Memory Management
-// ============================================================================
int fastnn_cuda_malloc(float** ptr, size_t size);
int fastnn_cuda_free(float* ptr);
int fastnn_cuda_memcpy_h2d(float* dst, const float* src, size_t size);
@@ -25,9 +19,6 @@ int fastnn_cuda_memcpy_d2h(float* dst, const float* src, size_t size);
int fastnn_cuda_memcpy_d2d(float* dst, const float* src, size_t size);
int fastnn_cuda_memset(float* ptr, int value, size_t size);
-// ============================================================================
-// Element-wise Operations
-// ============================================================================
int fastnn_cuda_add(const float* a, const float* b, float* out, size_t n);
int fastnn_cuda_sub(const float* a, const float* b, float* out, size_t n);
int fastnn_cuda_mul(const float* a, const float* b, float* out, size_t n);
@@ -42,9 +33,6 @@ int fastnn_cuda_exp(const float* a, float* out, size_t n);
int fastnn_cuda_log(const float* a, float* out, size_t n);
int fastnn_cuda_clamp(const float* a, float min_val, float max_val, float* out, size_t n);
-// ============================================================================
-// Activation Functions
-// ============================================================================
int fastnn_cuda_relu(const float* input, float* output, size_t n);
int fastnn_cuda_relu_backward(const float* grad_output, const float* input, float* grad_input, size_t n);
int fastnn_cuda_leaky_relu(const float* input, float negative_slope, float* output, size_t n);
@@ -58,16 +46,10 @@ int fastnn_cuda_gelu_backward(const float* grad_output, const float* input, floa
int fastnn_cuda_silu(const float* input, float* output, size_t n);
int fastnn_cuda_silu_backward(const float* grad_output, const float* input, float* grad_input, size_t n);
-// ============================================================================
-// Softmax
-// ============================================================================
int fastnn_cuda_softmax(const float* input, float* output, int batch_size, int num_classes);
int fastnn_cuda_softmax_backward(const float* grad_output, const float* output, float* grad_input, int batch_size, int num_classes);
int fastnn_cuda_log_softmax(const float* input, float* output, int batch_size, int num_classes);
-// ============================================================================
-// Matrix Operations (using custom kernels + cuBLAS)
-// ============================================================================
int fastnn_cuda_matmul(const float* a, const float* b, float* c,
int m, int n, int k,
int lda, int ldb, int ldc,
@@ -83,9 +65,6 @@ int fastnn_cuda_transpose(const float* input, float* output, int rows, int cols)
int fastnn_cuda_transpose_batched(const float* input, float* output, int batch_size, int rows, int cols);
int fastnn_cuda_permute_nd(const float* input, float* output, const int* out_strides, const int* in_strides, const int* perm, int ndim, int numel);
-// ============================================================================
-// Reduction Operations
-// ============================================================================
int fastnn_cuda_sum(const float* input, float* output, size_t n);
int fastnn_cuda_sum_axis(const float* input, float* output, int* shape, int ndim, int axis, int total_elements);
int fastnn_cuda_mean(const float* input, float* output, size_t n);
@@ -94,9 +73,6 @@ int fastnn_cuda_min(const float* input, float* output, size_t n);
int fastnn_cuda_argmax(const float* input, int* output, size_t n);
int fastnn_cuda_argmax_axis(const float* input, int* output, int outer, int axis_size, int inner);
-// ============================================================================
-// Convolution Operations
-// ============================================================================
int fastnn_cuda_conv2d_forward(
const float* input, const float* weight, const float* bias, float* output,
int batch_size, int in_channels, int out_channels,
@@ -121,9 +97,6 @@ int fastnn_cuda_conv2d_backward_weight(
int stride_h, int stride_w,
int pad_h, int pad_w);
-// ============================================================================
-// Pooling Operations
-// ============================================================================
int fastnn_cuda_max_pool2d(
const float* input, float* output, int* indices,
int batch_size, int channels,
@@ -152,9 +125,6 @@ int fastnn_cuda_adaptive_avg_pool2d(
int input_h, int input_w,
int output_h, int output_w);
-// ============================================================================
-// Normalization
-// ============================================================================
int fastnn_cuda_batch_norm_forward(
const float* input, const float* gamma, const float* beta,
float* running_mean, float* running_var,
@@ -184,28 +154,16 @@ int fastnn_cuda_rms_norm_forward(
float* output, float* rms,
int batch_size, int normalized_size, float epsilon);
-// ============================================================================
-// Dropout
-// ============================================================================
int fastnn_cuda_dropout_forward(const float* input, float* output, float* mask, float p, size_t n, unsigned long long seed);
int fastnn_cuda_dropout_backward(const float* grad_output, const float* mask, float* grad_input, float p, size_t n);
-// ============================================================================
-// Embedding
-// ============================================================================
int fastnn_cuda_embedding_forward(const int* indices, const float* weight, float* output, int num_indices, int embedding_dim);
int fastnn_cuda_embedding_backward(const int* indices, const float* grad_output, float* grad_weight, int num_indices, int embedding_dim, int num_embeddings);
-// ============================================================================
-// Loss Functions
-// ============================================================================
int fastnn_cuda_cross_entropy_loss(const float* log_probs, const int* targets, float* loss, float* grad, int batch_size, int num_classes);
int fastnn_cuda_mse_loss(const float* predictions, const float* targets, float* loss, float* grad, size_t n);
int fastnn_cuda_binary_cross_entropy(const float* predictions, const float* targets, float* loss, float* grad, size_t n);
-// ============================================================================
-// Attention / Transformer Operations
-// ============================================================================
int fastnn_cuda_scaled_dot_product_attention(
const float* query, const float* key, const float* value,
float* output, float* attn_weights,
@@ -213,9 +171,6 @@ int fastnn_cuda_scaled_dot_product_attention(
int batch_size, int num_heads, int seq_len_q, int seq_len_k, int head_dim,
float scale, int causal);
-// ============================================================================
-// Optimizer Kernels
-// ============================================================================
int fastnn_cuda_sgd_step(float* params, const float* grads, float* velocity,
float lr, float momentum, float weight_decay, float dampening,
int nesterov, size_t n);
@@ -226,9 +181,6 @@ int fastnn_cuda_adam_step(float* params, const float* grads,
float weight_decay, int step, int amsgrad, float* v_max,
size_t n);
-// ============================================================================
-// Utility Kernels
-// ============================================================================
int fastnn_cuda_fill(float* data, float value, size_t n);
int fastnn_cuda_copy(const float* src, float* dst, size_t n);
int fastnn_cuda_arange(float* output, float start, float step, size_t n);
diff --git a/cuda/kernels.cu b/cuda/kernels.cu
index 1af5c5f..497bc86 100644
--- a/cuda/kernels.cu
+++ b/cuda/kernels.cu
@@ -7,12 +7,11 @@
#include
#include
-// ============================================================================
-// Constants & Helpers
-// ============================================================================
#define BLOCK_SIZE 256
#define TILE_SIZE 32
#define WARP_SIZE 32
+#define ARGMAX_MAX_THREADS 1024
+#define PERMUTE_MAX_DIMS 8
#define CUDA_CHECK(call) do { \
cudaError_t err = (call); \
@@ -29,9 +28,6 @@ static curandGenerator_t g_curand_gen = nullptr;
static inline int div_ceil(int a, int b) { return (a + b - 1) / b; }
-// ============================================================================
-// Device Management
-// ============================================================================
extern "C" int fastnn_cuda_init(int device_id) {
CUDA_CHECK(cudaSetDevice(device_id));
@@ -63,9 +59,6 @@ extern "C" void fastnn_cuda_get_memory_info(size_t* free_mem, size_t* total_mem)
cudaMemGetInfo(free_mem, total_mem);
}
-// ============================================================================
-// Memory Management
-// ============================================================================
extern "C" int fastnn_cuda_malloc(float** ptr, size_t size) {
CUDA_CHECK(cudaMalloc(ptr, size));
return 0;
@@ -96,9 +89,6 @@ extern "C" int fastnn_cuda_memset(float* ptr, int value, size_t size) {
return 0;
}
-// ============================================================================
-// Element-wise Kernels
-// ============================================================================
__global__ void kernel_add(const float* a, const float* b, float* out, size_t n) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) out[idx] = a[idx] + b[idx];
@@ -164,13 +154,9 @@ __global__ void kernel_clamp(const float* a, float min_val, float max_val, float
if (idx < n) out[idx] = fminf(fmaxf(a[idx], min_val), max_val);
}
-// Threads for a one-block-per-row reduction over `n` values.
-//
-// The tree reductions below halve blockDim.x each round, so a thread count that
-// is not a power of two silently drops the tail: with 5 threads, stride 2 folds
-// 2->0 and 3->1 and thread 4 is never read. Round up instead. Threads past `n`
-// skip their strided loop entirely and so contribute the identity (0 for a sum,
-// -FLT_MAX for a max), which is exactly what the reduction needs.
+// The tree reductions halve blockDim.x each round, so a non-power-of-two thread
+// count drops the tail (5 threads: stride 2 folds 2->0 and 3->1, thread 4 is
+// never read). Round up; extra threads skip their loop and add the identity.
__host__ inline int reduction_threads(int n) {
int threads = 1;
while (threads < n && threads < BLOCK_SIZE) threads <<= 1;
@@ -199,9 +185,6 @@ extern "C" int fastnn_cuda_exp(const float* a, float* out, size_t n) { LAUNCH_EL
extern "C" int fastnn_cuda_log(const float* a, float* out, size_t n) { LAUNCH_ELEMENTWISE(kernel_log, a, out, n); }
extern "C" int fastnn_cuda_clamp(const float* a, float min_v, float max_v, float* out, size_t n) { LAUNCH_ELEMENTWISE(kernel_clamp, a, min_v, max_v, out, n); }
-// ============================================================================
-// Activation Kernels
-// ============================================================================
__global__ void kernel_relu(const float* input, float* output, size_t n) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) output[idx] = fmaxf(input[idx], 0.0f);
@@ -245,7 +228,6 @@ __global__ void kernel_tanh_backward(const float* grad_output, const float* outp
if (idx < n) grad_input[idx] = grad_output[idx] * (1.0f - output[idx] * output[idx]);
}
-// GELU: 0.5 * x * (1 + erf(x / sqrt(2)))
__global__ void kernel_gelu(const float* input, float* output, size_t n) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
@@ -264,7 +246,6 @@ __global__ void kernel_gelu_backward(const float* grad_output, const float* inpu
}
}
-// SiLU (Swish): x * sigmoid(x)
__global__ void kernel_silu(const float* input, float* output, size_t n) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
@@ -296,9 +277,6 @@ extern "C" int fastnn_cuda_gelu_backward(const float* go, const float* in, float
extern "C" int fastnn_cuda_silu(const float* in, float* out, size_t n) { LAUNCH_ELEMENTWISE(kernel_silu, in, out, n); }
extern "C" int fastnn_cuda_silu_backward(const float* go, const float* in, float* gi, size_t n) { LAUNCH_ELEMENTWISE(kernel_silu_backward, go, in, gi, n); }
-// ============================================================================
-// Softmax (numerically stable, shared memory)
-// ============================================================================
__global__ void kernel_softmax(const float* input, float* output, int batch_size, int num_classes) {
int batch_idx = blockIdx.x;
if (batch_idx >= batch_size) return;
@@ -307,7 +285,7 @@ __global__ void kernel_softmax(const float* input, float* output, int batch_size
const float* in_row = input + batch_idx * num_classes;
float* out_row = output + batch_idx * num_classes;
- // Find max for numerical stability
+ // subtract the row max so expf cannot overflow
float thread_max = -FLT_MAX;
for (int i = threadIdx.x; i < num_classes; i += blockDim.x) {
thread_max = fmaxf(thread_max, in_row[i]);
@@ -315,7 +293,6 @@ __global__ void kernel_softmax(const float* input, float* output, int batch_size
shared[threadIdx.x] = thread_max;
__syncthreads();
- // Parallel reduction for max
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (threadIdx.x < stride)
shared[threadIdx.x] = fmaxf(shared[threadIdx.x], shared[threadIdx.x + stride]);
@@ -324,7 +301,6 @@ __global__ void kernel_softmax(const float* input, float* output, int batch_size
float max_val = shared[0];
__syncthreads();
- // Compute exp(x - max) and sum
float thread_sum = 0.0f;
for (int i = threadIdx.x; i < num_classes; i += blockDim.x) {
float val = expf(in_row[i] - max_val);
@@ -334,7 +310,6 @@ __global__ void kernel_softmax(const float* input, float* output, int batch_size
shared[threadIdx.x] = thread_sum;
__syncthreads();
- // Parallel reduction for sum
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (threadIdx.x < stride)
shared[threadIdx.x] += shared[threadIdx.x + stride];
@@ -343,7 +318,6 @@ __global__ void kernel_softmax(const float* input, float* output, int batch_size
float sum_val = shared[0];
__syncthreads();
- // Normalize
for (int i = threadIdx.x; i < num_classes; i += blockDim.x) {
out_row[i] /= sum_val;
}
@@ -359,7 +333,6 @@ __global__ void kernel_softmax_backward(const float* grad_output, const float* o
const float* out_row = output + batch_idx * num_classes;
float* gi_row = grad_input + batch_idx * num_classes;
- // Compute dot product: sum(grad_output * output)
float thread_dot = 0.0f;
for (int i = threadIdx.x; i < num_classes; i += blockDim.x) {
thread_dot += go_row[i] * out_row[i];
@@ -375,7 +348,6 @@ __global__ void kernel_softmax_backward(const float* grad_output, const float* o
float dot = shared[0];
__syncthreads();
- // grad_input = output * (grad_output - dot)
for (int i = threadIdx.x; i < num_classes; i += blockDim.x) {
gi_row[i] = out_row[i] * (go_row[i] - dot);
}
@@ -438,14 +410,11 @@ extern "C" int fastnn_cuda_log_softmax(const float* input, float* output, int bs
return 0;
}
-// ============================================================================
-// Matrix Operations (cuBLAS for GEMM)
-// ============================================================================
extern "C" int fastnn_cuda_matmul(const float* a, const float* b, float* c,
int m, int n, int k,
int lda, int ldb, int ldc,
float alpha, float beta) {
- // cuBLAS uses column-major, so we compute C^T = B^T * A^T
+ // cuBLAS is column-major, so compute C^T = B^T * A^T
CUBLAS_CHECK(cublasSgemm(g_cublas_handle,
CUBLAS_OP_N, CUBLAS_OP_N,
n, m, k,
@@ -457,15 +426,10 @@ extern "C" int fastnn_cuda_matmul(const float* a, const float* b, float* c,
return 0;
}
-// ─── Transpose-fused GEMM variants ─────────────────────────────────────────
-// All matrices are row-major. cuBLAS is column-major so we use the standard
-// "C^T = B^T × A^T" trick: pass arguments in reversed order with swapped dims.
-//
-// C (m×n) = A (m×k) × B^T (k×n) where B is stored (n×k) row-major.
+// C (m×n) = A (m×k) × B^T, B stored (n×k)
extern "C" int fastnn_cuda_matmul_nt(const float* a, const float* b, float* c,
int m, int n, int k) {
float alpha = 1.0f, beta = 0.0f;
- // C^T (n×m) = B (n×k) × A^T (k×m)
CUBLAS_CHECK(cublasSgemm(g_cublas_handle,
CUBLAS_OP_T, CUBLAS_OP_N,
n, m, k,
@@ -473,11 +437,10 @@ extern "C" int fastnn_cuda_matmul_nt(const float* a, const float* b, float* c,
return 0;
}
-// C (m×n) = A^T (m×k) × B (k×n) where A is stored (k×m) row-major.
+// C (m×n) = A^T × B (k×n), A stored (k×m)
extern "C" int fastnn_cuda_matmul_tn(const float* a, const float* b, float* c,
int m, int n, int k) {
float alpha = 1.0f, beta = 0.0f;
- // C^T (n×m) = B^T (n×k) × A (k×m) — A (k×m) row-major, lda=m
CUBLAS_CHECK(cublasSgemm(g_cublas_handle,
CUBLAS_OP_N, CUBLAS_OP_T,
n, m, k,
@@ -485,7 +448,7 @@ extern "C" int fastnn_cuda_matmul_tn(const float* a, const float* b, float* c,
return 0;
}
-// Batched C (batch×m×n) = A (batch×m×k) × B^T (batch×k×n) where B is (batch×n×k).
+// Batched C = A × B^T, B stored (batch×n×k)
extern "C" int fastnn_cuda_matmul_batched_nt(const float* a, const float* b, float* c,
int m, int n, int k, int batch_size) {
float alpha = 1.0f, beta = 0.0f;
@@ -502,7 +465,7 @@ extern "C" int fastnn_cuda_matmul_batched_nt(const float* a, const float* b, flo
return 0;
}
-// Batched C (batch×m×n) = A^T (batch×m×k) × B (batch×k×n) where A is (batch×k×m).
+// Batched C = A^T × B, A stored (batch×k×m)
extern "C" int fastnn_cuda_matmul_batched_tn(const float* a, const float* b, float* c,
int m, int n, int k, int batch_size) {
float alpha = 1.0f, beta = 0.0f;
@@ -539,7 +502,7 @@ extern "C" int fastnn_cuda_matmul_batched(const float* a, const float* b, float*
}
__global__ void kernel_transpose(const float* input, float* output, int rows, int cols) {
- __shared__ float tile[TILE_SIZE][TILE_SIZE + 1]; // +1 to avoid bank conflicts
+ __shared__ float tile[TILE_SIZE][TILE_SIZE + 1]; // +1 pad avoids bank conflicts
int x = blockIdx.x * TILE_SIZE + threadIdx.x;
int y = blockIdx.y * TILE_SIZE + threadIdx.y;
@@ -594,13 +557,10 @@ extern "C" int fastnn_cuda_transpose_batched(const float* input, float* output,
return 0;
}
-// ============================================================================
-// N-D Permute
-// ============================================================================
struct PermuteNDParams {
- int out_strides[8];
- int in_strides[8];
- int perm[8];
+ int out_strides[PERMUTE_MAX_DIMS];
+ int in_strides[PERMUTE_MAX_DIMS];
+ int perm[PERMUTE_MAX_DIMS];
int ndim;
};
@@ -630,7 +590,7 @@ extern "C" int fastnn_cuda_permute_nd(
{
PermuteNDParams params;
params.ndim = ndim;
- for (int i = 0; i < ndim && i < 8; i++) {
+ for (int i = 0; i < ndim && i < PERMUTE_MAX_DIMS; i++) {
params.out_strides[i] = out_strides[i];
params.in_strides[i] = in_strides[i];
params.perm[i] = perm[i];
@@ -641,9 +601,6 @@ extern "C" int fastnn_cuda_permute_nd(
return 0;
}
-// ============================================================================
-// Reduction Operations
-// ============================================================================
__global__ void kernel_sum(const float* input, float* output, size_t n) {
extern __shared__ float sdata[];
unsigned int tid = threadIdx.x;
@@ -660,7 +617,6 @@ __global__ void kernel_sum(const float* input, float* output, size_t n) {
__syncthreads();
}
- // Warp-level reduction (no sync needed)
if (tid < WARP_SIZE) {
volatile float* smem = sdata;
if (blockDim.x >= 64) smem[tid] += smem[tid + 32];
@@ -685,13 +641,11 @@ extern "C" int fastnn_cuda_sum(const float* input, float* output, size_t n) {
__global__ void kernel_sum_axis(const float* input, float* output, int* shape, int ndim, int axis, int total_elements) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
- // Calculate output size
int axis_size = shape[axis];
int output_size = total_elements / axis_size;
if (idx >= output_size) return;
- // Calculate the multi-dim index in the output
int inner_size = 1;
for (int i = axis + 1; i < ndim; i++) inner_size *= shape[i];
@@ -707,19 +661,12 @@ __global__ void kernel_sum_axis(const float* input, float* output, int* shape, i
}
extern "C" int fastnn_cuda_sum_axis(const float* input, float* output, int* shape, int ndim, int axis, int total_elements) {
- int axis_size = 1;
- // Read axis size from device memory would be complex, so we calculate output size
- int output_size = total_elements; // Will be divided by axis_size in kernel
- // We need to copy shape to device
+ // shape is host memory, the kernel needs its own device copy
int* d_shape;
CUDA_CHECK(cudaMalloc(&d_shape, ndim * sizeof(int)));
CUDA_CHECK(cudaMemcpy(d_shape, shape, ndim * sizeof(int), cudaMemcpyHostToDevice));
- // Calculate output size on host
- int h_axis_size;
- // shape is host memory in this interface
- h_axis_size = shape[axis];
- output_size = total_elements / h_axis_size;
+ int output_size = total_elements / shape[axis];
int blocks = div_ceil(output_size, BLOCK_SIZE);
kernel_sum_axis<<>>(input, output, d_shape, ndim, axis, total_elements);
@@ -751,7 +698,6 @@ __global__ void kernel_max(const float* input, float* output, size_t n) {
}
if (tid == 0) {
- // Atomic max for float using atomicCAS
float old = *output;
while (sdata[0] > old) {
float assumed = old;
@@ -771,7 +717,6 @@ extern "C" int fastnn_cuda_max(const float* input, float* output, size_t n) {
}
extern "C" int fastnn_cuda_min(const float* input, float* output, size_t n) {
- // Negate, find max, negate result
float* temp;
CUDA_CHECK(cudaMalloc(&temp, n * sizeof(float)));
fastnn_cuda_neg(input, temp, n);
@@ -802,15 +747,14 @@ __global__ void kernel_argmax(const float* input, int* output, size_t n) {
}
if (tid == 0) {
- // Simple atomic compare-and-swap for global argmax
- // For single-block case, just write directly
output[0] = sidx[0];
}
}
extern "C" int fastnn_cuda_argmax(const float* input, int* output, size_t n) {
- // Simple single-pass for now
- kernel_argmax<<<1, min((int)n, 1024), 1024 * (sizeof(float) + sizeof(int))>>>(input, output, n);
+ // One block, one element per thread: only right for n <= ARGMAX_MAX_THREADS
+ // and a power-of-two n.
+ kernel_argmax<<<1, min((int)n, ARGMAX_MAX_THREADS), ARGMAX_MAX_THREADS * (sizeof(float) + sizeof(int))>>>(input, output, n);
CUDA_CHECK(cudaGetLastError());
return 0;
}
@@ -843,9 +787,6 @@ extern "C" int fastnn_cuda_argmax_axis(const float* input, int* output, int oute
return 0;
}
-// ============================================================================
-// Convolution (im2col + GEMM approach)
-// ============================================================================
__global__ void kernel_im2col(
const float* input, float* col,
int channels, int height, int width,
@@ -925,16 +866,11 @@ extern "C" int fastnn_cuda_conv2d_forward(
const float* input_b = input + b * in_channels * input_h * input_w;
float* output_b = output + b * out_channels * out_h * out_w;
- // im2col
int blocks = div_ceil(col_size, BLOCK_SIZE);
kernel_im2col<<>>(
input_b, col, in_channels, input_h, input_w,
kernel_h, kernel_w, stride_h, stride_w, pad_h, pad_w, out_h, out_w);
- // GEMM: output = weight * col
- // weight: [out_channels, in_channels * kernel_h * kernel_w]
- // col: [in_channels * kernel_h * kernel_w, out_h * out_w]
- // output: [out_channels, out_h * out_w]
float alpha = 1.0f, beta = 0.0f;
int M = out_channels;
int N = out_h * out_w;
@@ -943,7 +879,6 @@ extern "C" int fastnn_cuda_conv2d_forward(
CUBLAS_CHECK(cublasSgemm(g_cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N,
N, M, K, &alpha, col, N, weight, K, &beta, output_b, N));
- // Add bias if provided
if (bias != nullptr) {
for (int oc = 0; oc < out_channels; oc++) {
float bias_val;
@@ -977,7 +912,6 @@ extern "C" int fastnn_cuda_conv2d_backward_data(
const float* go_b = grad_output + b * out_channels * output_h * output_w;
float* gi_b = grad_input + b * in_channels * input_h * input_w;
- // col = weight^T * grad_output
float alpha = 1.0f, beta = 0.0f;
int M = in_channels * kernel_h * kernel_w;
int N = output_h * output_w;
@@ -986,7 +920,6 @@ extern "C" int fastnn_cuda_conv2d_backward_data(
CUBLAS_CHECK(cublasSgemm(g_cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T,
N, M, K, &alpha, go_b, N, weight, M, &beta, col, N));
- // col2im
int total = in_channels * input_h * input_w;
int blocks = div_ceil(total, BLOCK_SIZE);
kernel_col2im<<>>(
@@ -1017,13 +950,11 @@ extern "C" int fastnn_cuda_conv2d_backward_weight(
const float* input_b = input + b * in_channels * input_h * input_w;
const float* go_b = grad_output + b * out_channels * output_h * output_w;
- // im2col on input
int blocks = div_ceil(col_size, BLOCK_SIZE);
kernel_im2col<<>>(
input_b, col, in_channels, input_h, input_w,
kernel_h, kernel_w, stride_h, stride_w, pad_h, pad_w, output_h, output_w);
- // grad_weight += grad_output * col^T
float alpha = 1.0f, beta = 1.0f;
int M = out_channels;
int N = in_channels * kernel_h * kernel_w;
@@ -1038,9 +969,6 @@ extern "C" int fastnn_cuda_conv2d_backward_weight(
return 0;
}
-// ============================================================================
-// Pooling Operations
-// ============================================================================
__global__ void kernel_max_pool2d(
const float* input, float* output, int* indices,
int batch_size, int channels,
@@ -1234,9 +1162,6 @@ extern "C" int fastnn_cuda_adaptive_avg_pool2d(
return 0;
}
-// ============================================================================
-// Batch Normalization
-// ============================================================================
__global__ void kernel_batch_norm_forward(
const float* input, const float* gamma, const float* beta,
float* running_mean, float* running_var,
@@ -1250,14 +1175,12 @@ __global__ void kernel_batch_norm_forward(
int n = batch_size * spatial_size;
if (training) {
- // Compute mean
float mean = 0.0f;
for (int i = threadIdx.x; i < n; i += blockDim.x) {
int b = i / spatial_size;
int s = i % spatial_size;
mean += input[(b * channels + c) * spatial_size + s];
}
- // Warp reduction
__shared__ float smean[BLOCK_SIZE];
smean[threadIdx.x] = mean;
__syncthreads();
@@ -1267,7 +1190,6 @@ __global__ void kernel_batch_norm_forward(
}
mean = smean[0] / (float)n;
- // Compute variance
float var = 0.0f;
for (int i = threadIdx.x; i < n; i += blockDim.x) {
int b = i / spatial_size;
@@ -1294,7 +1216,6 @@ __global__ void kernel_batch_norm_forward(
}
__syncthreads();
- // Normalize
for (int i = threadIdx.x; i < n; i += blockDim.x) {
int b = i / spatial_size;
int s = i % spatial_size;
@@ -1343,7 +1264,6 @@ __global__ void kernel_batch_norm_backward(
float inv_var = save_inv_var[c];
float g = gamma[c];
- // Compute grad_gamma and grad_beta
float dg = 0.0f, db = 0.0f;
for (int i = threadIdx.x; i < n; i += blockDim.x) {
int b = i / spatial_size;
@@ -1374,7 +1294,6 @@ __global__ void kernel_batch_norm_backward(
}
__syncthreads();
- // Compute grad_input
float inv_n = 1.0f / (float)n;
for (int i = threadIdx.x; i < n; i += blockDim.x) {
int b = i / spatial_size;
@@ -1399,9 +1318,6 @@ extern "C" int fastnn_cuda_batch_norm_backward(
return 0;
}
-// ============================================================================
-// Layer Normalization
-// ============================================================================
__global__ void kernel_layer_norm_forward(
const float* input, const float* gamma, const float* beta,
float* output, float* mean_out, float* inv_var_out,
@@ -1414,7 +1330,6 @@ __global__ void kernel_layer_norm_forward(
const float* in_row = input + b * normalized_size;
float* out_row = output + b * normalized_size;
- // Compute mean
float thread_sum = 0.0f;
for (int i = threadIdx.x; i < normalized_size; i += blockDim.x)
thread_sum += in_row[i];
@@ -1427,7 +1342,6 @@ __global__ void kernel_layer_norm_forward(
float mean = shared[0] / (float)normalized_size;
__syncthreads();
- // Compute variance
float thread_var = 0.0f;
for (int i = threadIdx.x; i < normalized_size; i += blockDim.x) {
float diff = in_row[i] - mean;
@@ -1447,7 +1361,6 @@ __global__ void kernel_layer_norm_forward(
}
__syncthreads();
- // Normalize
for (int i = threadIdx.x; i < normalized_size; i += blockDim.x) {
float normalized = (in_row[i] - mean) * inv_var;
out_row[i] = gamma[i] * normalized + beta[i];
@@ -1482,7 +1395,6 @@ __global__ void kernel_layer_norm_backward(
float m = mean[b];
float iv = inv_var[b];
- // Compute partial sums
float ds = 0.0f, db = 0.0f;
for (int i = threadIdx.x; i < normalized_size; i += blockDim.x) {
float xhat = (in_row[i] - m) * iv;
@@ -1492,7 +1404,6 @@ __global__ void kernel_layer_norm_backward(
atomicAdd(&grad_beta[i], go_row[i]);
}
- // Reduce ds and db
float* shared2 = shared + blockDim.x;
shared[threadIdx.x] = ds;
shared2[threadIdx.x] = db;
@@ -1531,7 +1442,6 @@ extern "C" int fastnn_cuda_layer_norm_backward(
return 0;
}
-// RMS Norm
__global__ void kernel_rms_norm_forward(
const float* input, const float* gamma,
float* output, float* rms_out,
@@ -1573,9 +1483,6 @@ extern "C" int fastnn_cuda_rms_norm_forward(
return 0;
}
-// ============================================================================
-// Dropout
-// ============================================================================
__global__ void kernel_dropout_forward(const float* input, float* output, float* mask, float p, size_t n, unsigned long long seed) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= n) return;
@@ -1609,9 +1516,6 @@ extern "C" int fastnn_cuda_dropout_backward(const float* grad_output, const floa
LAUNCH_ELEMENTWISE(kernel_dropout_backward, grad_output, mask, grad_input, p, n);
}
-// ============================================================================
-// Embedding
-// ============================================================================
__global__ void kernel_embedding_forward(const int* indices, const float* weight, float* output, int num_indices, int embedding_dim) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= num_indices * embedding_dim) return;
@@ -1651,9 +1555,6 @@ extern "C" int fastnn_cuda_embedding_backward(const int* indices, const float* g
return 0;
}
-// ============================================================================
-// Loss Functions
-// ============================================================================
__global__ void kernel_cross_entropy_loss(const float* log_probs, const int* targets, float* loss, float* grad,
int batch_size, int num_classes) {
int b = blockIdx.x * blockDim.x + threadIdx.x;
@@ -1663,7 +1564,6 @@ __global__ void kernel_cross_entropy_loss(const float* log_probs, const int* tar
float nll = -log_probs[b * num_classes + target];
atomicAdd(loss, nll / (float)batch_size);
- // Gradient: softmax output - one_hot (for log_softmax input)
for (int c = 0; c < num_classes; c++) {
float prob = expf(log_probs[b * num_classes + c]);
grad[b * num_classes + c] = (prob - (c == target ? 1.0f : 0.0f)) / (float)batch_size;
@@ -1714,9 +1614,6 @@ extern "C" int fastnn_cuda_binary_cross_entropy(const float* predictions, const
return 0;
}
-// ============================================================================
-// Scaled Dot-Product Attention
-// ============================================================================
extern "C" int fastnn_cuda_scaled_dot_product_attention(
const float* query, const float* key, const float* value,
float* output, float* attn_weights,
@@ -1726,16 +1623,14 @@ extern "C" int fastnn_cuda_scaled_dot_product_attention(
{
int bh = batch_size * num_heads;
- // QK^T: [bh, seq_len_q, head_dim] x [bh, head_dim, seq_len_k] -> [bh, seq_len_q, seq_len_k]
float alpha = scale;
float beta = 0.0f;
- // We need to transpose K: [bh, seq_len_k, head_dim] -> [bh, head_dim, seq_len_k]
float* key_t;
CUDA_CHECK(cudaMalloc(&key_t, (size_t)bh * head_dim * seq_len_k * sizeof(float)));
fastnn_cuda_transpose_batched(key, key_t, bh, seq_len_k, head_dim);
- // Batched matmul: scores = Q * K^T * scale
+ // scores = scale * Q K^T
CUBLAS_CHECK(cublasSgemmStridedBatched(g_cublas_handle,
CUBLAS_OP_N, CUBLAS_OP_N,
seq_len_k, seq_len_q, head_dim,
@@ -1748,30 +1643,9 @@ extern "C" int fastnn_cuda_scaled_dot_product_attention(
cudaFree(key_t);
- // Apply causal mask if needed
- if (causal) {
- // Simple causal mask: set upper triangular to -inf
- int total = bh * seq_len_q * seq_len_k;
- int blocks = div_ceil(total, BLOCK_SIZE);
- // We'll use a lambda-style kernel
- auto causal_mask = [=] __device__ (int idx) {
- int k_pos = idx % seq_len_k;
- int q_pos = (idx / seq_len_k) % seq_len_q;
- if (k_pos > q_pos) return -1e9f;
- return 0.0f;
- };
- // Use a kernel for this
- // (Inline kernel for causal masking)
- }
-
- // Softmax over last dimension
- for (int i = 0; i < bh * seq_len_q; i++) {
- // Apply softmax per row (seq_len_q rows, each of length seq_len_k)
- }
- // Use our existing softmax
+ // `causal` and `mask` are not applied yet.
fastnn_cuda_softmax(attn_weights, attn_weights, bh * seq_len_q, seq_len_k);
- // Output = attn_weights * V
alpha = 1.0f;
beta = 0.0f;
CUBLAS_CHECK(cublasSgemmStridedBatched(g_cublas_handle,
@@ -1788,9 +1662,6 @@ extern "C" int fastnn_cuda_scaled_dot_product_attention(
return 0;
}
-// ============================================================================
-// Optimizer Kernels
-// ============================================================================
__global__ void kernel_sgd_step(float* params, const float* grads, float* velocity,
float lr, float momentum, float weight_decay, float dampening,
int nesterov, size_t n) {
@@ -1830,16 +1701,14 @@ __global__ void kernel_adam_step(float* params, const float* grads,
float g = grads[idx];
- // AdamW weight decay (decoupled)
+ // decay is applied to the weights directly, not folded into the gradient
if (weight_decay != 0.0f) {
params[idx] -= lr * weight_decay * params[idx];
}
- // Update biased first and second moment estimates
m[idx] = beta1 * m[idx] + (1.0f - beta1) * g;
v[idx] = beta2 * v[idx] + (1.0f - beta2) * g * g;
- // Bias correction
float m_hat = m[idx] / (1.0f - powf(beta1, (float)step));
float v_hat = v[idx] / (1.0f - powf(beta2, (float)step));
@@ -1862,9 +1731,6 @@ extern "C" int fastnn_cuda_adam_step(float* params, const float* grads,
return 0;
}
-// ============================================================================
-// Utility Kernels
-// ============================================================================
__global__ void kernel_fill(float* data, float value, size_t n) {
size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) data[idx] = value;
@@ -1924,9 +1790,6 @@ extern "C" int fastnn_cuda_gather(const float* input, const int* indices, float*
return 0;
}
-// ============================================================================
-// Convolution lowering: im2col / col2im
-// ============================================================================
__global__ void kernel_im2col(const float* input, float* cols,
int c, int h, int w,
int kh, int kw, int sh, int sw, int ph, int pw,
@@ -1963,9 +1826,7 @@ extern "C" int fastnn_cuda_im2col(const float* input, float* cols,
return 0;
}
-// One thread per input pixel, summing every window that covered it. Walking
-// the windows from the pixel's side needs no atomics, unlike scattering from
-// the columns' side.
+// One thread per input pixel gathers its windows, so no atomics are needed.
__global__ void kernel_col2im(const float* cols, float* image,
int c, int h, int w,
int kh, int kw, int sh, int sw, int ph, int pw,
diff --git a/cuda/stubs.c b/cuda/stubs.c
index 4fe5000..7f460fc 100644
--- a/cuda/stubs.c
+++ b/cuda/stubs.c
@@ -1,8 +1,4 @@
-/* CPU-only build stubs for CUDA kernel symbols.
- * Compiled when the `cuda` feature is disabled. These functions are never
- * called at runtime (all tensor storage is CPU in that configuration), but
- * the Rust FFI declarations require symbols to be present at link time.
- */
+/* Link-time stand-ins for the kernel symbols when the `cuda` feature is off. Never called. */
#include
#define STUB return -1
diff --git a/examples/char_lm.rs b/examples/char_lm.rs
index 3c1e89b..04f3cd4 100644
--- a/examples/char_lm.rs
+++ b/examples/char_lm.rs
@@ -1,28 +1,24 @@
-//! A character-level language model: a small GPT trained from scratch.
+//! Character-level GPT trained from scratch.
//!
-//! cargo run --example char_lm --release
-//! cargo run --example char_lm --release -- corpus.txt
+//! cargo run --example char_lm --release [-- corpus.txt]
//!
-//! With no argument it trains on a short embedded excerpt, which is enough to
-//! see the loss fall and the samples turn into English-shaped nonsense. For text
-//! worth reading, point it at a real corpus:
-//!
-//! curl -o shakespeare.txt \
-//! https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
-//!
-//! It checkpoints as it goes and resumes from `char_lm.fdl` if one is there, so
-//! stopping it with Ctrl-C costs at most a few hundred steps.
+//! Defaults to a small embedded excerpt. Checkpoints to `char_lm.fdl` and
+//! resumes from it on the next run.
use fastnn::nn::module::scoped;
use fastnn::prelude::*;
-/// Where an interrupted run is saved so it can be picked up again.
const CHECKPOINT: &str = "char_lm.fdl";
+const SEED: u64 = 1337;
+const DROPOUT: f32 = 0.1;
+const CLIP_NORM: f32 = 1.0;
+const SAMPLE_CHARS: usize = 400;
+const TEMPERATURE: f32 = 0.8;
const CORPUS: &str = include_str!("data/tiny_corpus.txt");
fn main() -> fastnn::Result<()> {
- manual_seed(1337);
+ manual_seed(SEED);
let text = match std::env::args().nth(1) {
Some(path) => std::fs::read_to_string(&path)?,
@@ -45,15 +41,12 @@ fn main() -> fastnn::Result<()> {
let mut opt = AdamW::new(model.parameters(), config.lr).betas(0.9, 0.95);
- // Resume if a previous run left a checkpoint, otherwise start from scratch.
- // Weights alone would not be enough: AdamW's moments and step count decide
- // how large the next update is.
+ // Resuming needs the optimizer state too, not just the weights.
let first_step = load_training(&model, &mut opt, CHECKPOINT).unwrap_or(0) as usize;
if first_step > 0 {
println!("resuming from step {first_step}\n");
}
- // Warmup first: Adam's variance estimate is unreliable for the first few
- // hundred steps, and a full-size update built on it can wreck the model.
+ // Adam's variance estimate is noisy early on, so ramp the lr up.
let schedule = Warmup::new(CosineAnnealing::new(config.lr, config.steps), config.warmup);
let params = model.parameters();
@@ -66,9 +59,7 @@ fn main() -> fastnn::Result<()> {
opt.zero_grad();
loss.backward();
- // Attention gradients spike occasionally; clipping keeps one bad batch
- // from undoing many good ones.
- clip_grad_norm(¶ms, 1.0);
+ clip_grad_norm(¶ms, CLIP_NORM);
opt.step();
if (step + 1) % config.report_every == 0 {
@@ -86,15 +77,14 @@ fn main() -> fastnn::Result<()> {
}
model.eval();
- println!("\n─── sample ───");
- println!("{}", generate(&model, &vocab, &config, 400, 0.8));
+ println!("\nsample:");
+ println!(
+ "{}",
+ generate(&model, &vocab, &config, SAMPLE_CHARS, TEMPERATURE)
+ );
Ok(())
}
-// ── Model ────────────────────────────────────────────────────────────────────
-
-/// Token embeddings plus learned positions, a causal transformer stack, and a
-/// linear head back to vocabulary logits.
struct CharGPT {
tokens: Embedding,
positions: Embedding,
@@ -113,17 +103,14 @@ impl CharGPT {
config.heads,
4 * config.width,
config.layers,
- 0.1,
+ DROPOUT,
),
head: Linear::no_bias(config.width, vocab),
context: config.context,
}
}
- /// `[batch][context]` token ids → `[batch·context, vocab]` logits.
- ///
- /// Flattened because that is the shape [`cross_entropy`] takes: every
- /// position in every sequence is one independent next-token prediction.
+ /// Returns `[batch * length, vocab]` logits, flattened for `cross_entropy`.
fn forward_ids(&self, ids: &[Vec]) -> Tensor {
let (batch, length) = (ids.len(), ids[0].len());
assert!(
@@ -173,9 +160,6 @@ impl Module for CharGPT {
}
}
-// ── Data ─────────────────────────────────────────────────────────────────────
-
-/// A character alphabet built from the corpus itself.
struct Vocab {
to_id: std::collections::HashMap,
to_char: Vec,
@@ -183,12 +167,11 @@ struct Vocab {
impl Vocab {
fn new(text: &str) -> Vocab {
- let mut to_char: Vec = text
+ let to_char: Vec = text
.chars()
.collect::>()
.into_iter()
.collect();
- to_char.shrink_to_fit();
Vocab {
to_id: to_char.iter().enumerate().map(|(i, &c)| (c, i)).collect(),
to_char,
@@ -210,8 +193,7 @@ impl Vocab {
}
}
-/// A batch of random windows: each input is `context` characters and each target
-/// is the same window shifted one step, so position `i` predicts character `i+1`.
+/// Random windows; targets are the inputs shifted by one character.
fn sample_batch(data: &[usize], config: &Config) -> (Vec>, Vec) {
let mut inputs = Vec::with_capacity(config.batch);
let mut targets = Vec::with_capacity(config.batch * config.context);
@@ -225,9 +207,6 @@ fn sample_batch(data: &[usize], config: &Config) -> (Vec>, Vec
(inputs, targets)
}
-// ── Generation ───────────────────────────────────────────────────────────────
-
-/// Sample `count` characters, feeding each one back in as context.
fn generate(
model: &CharGPT,
vocab: &Vocab,
@@ -242,7 +221,6 @@ fn generate(
let context = vec![ids[window..].to_vec()];
let logits = no_grad(|| model.forward_ids(&context));
- // Only the final position predicts the next character.
let last = logits.narrow(0, logits.dim(0) - 1, 1);
ids.push(sample(&last.div_scalar(temperature).softmax().to_vec()));
}
@@ -250,7 +228,6 @@ fn generate(
vocab.decode(&ids)
}
-/// Draw one index from a probability vector by inverse-CDF.
fn sample(probabilities: &[f32]) -> usize {
let mut draw = fastnn::rng::uniform(1, 0.0, 1.0)[0];
for (index, &p) in probabilities.iter().enumerate() {
@@ -262,9 +239,6 @@ fn sample(probabilities: &[f32]) -> usize {
probabilities.len() - 1
}
-// ── Configuration ────────────────────────────────────────────────────────────
-
-/// Model and schedule sizes, picked from how much text there is to learn from.
#[derive(Debug)]
struct Config {
context: usize,
@@ -281,8 +255,7 @@ struct Config {
impl Config {
fn for_corpus(chars: usize) -> Config {
- // A model much larger than its corpus memorises instead of generalising,
- // so capacity scales with the amount of text available.
+ // A model much bigger than its corpus just memorises it.
if chars > 200_000 {
Config {
context: 128,
diff --git a/examples/finetune.rs b/examples/finetune.rs
index 2b60bf7..fb5ae91 100644
--- a/examples/finetune.rs
+++ b/examples/finetune.rs
@@ -1,17 +1,17 @@
-//! Fine-tuning: pretrain on one task, freeze the backbone, retrain a new head
-//! on another.
+//! Pretrain on quadrant classification, freeze the backbone, train a new head
+//! for inside/outside the unit circle.
//!
//! cargo run --example finetune --release
-//!
-//! Task A classifies 2-D points by quadrant; task B by whether they fall inside
-//! the unit circle. The backbone's features transfer, so the second task trains
-//! a 34-parameter head instead of a whole network — and because frozen
-//! parameters hand out detached tensors, the backward pass never even reaches
-//! the backbone.
use fastnn::nn::module::scoped;
use fastnn::prelude::*;
+const SEED: u64 = 11;
+const SAMPLES: usize = 512;
+const HIDDEN: usize = 16;
+const LR: f32 = 1e-2;
+const STEPS: usize = 300;
+
struct Net {
backbone: Sequential,
head: Linear,
@@ -21,11 +21,11 @@ impl Net {
fn new(classes: usize) -> Net {
Net {
backbone: Sequential::new()
- .add(Linear::new(2, 16))
+ .add(Linear::new(2, HIDDEN))
.add(Tanh)
- .add(Linear::new(16, 16))
+ .add(Linear::new(HIDDEN, HIDDEN))
.add(Tanh),
- head: Linear::new(16, classes),
+ head: Linear::new(HIDDEN, classes),
}
}
}
@@ -41,7 +41,6 @@ impl Module for Net {
}
}
-/// 2-D points and their labels under `rule`.
fn dataset(n: usize, rule: impl Fn(f32, f32) -> usize) -> (Tensor, Vec) {
let coords = fastnn::rng::uniform(n * 2, -1.5, 1.5);
let labels = coords.chunks(2).map(|p| rule(p[0], p[1])).collect();
@@ -55,8 +54,7 @@ fn accuracy(model: &impl Module, inputs: &Tensor, labels: &[usize]) -> f32 {
}
fn fit(model: &impl Module, inputs: &Tensor, labels: &[usize], steps: usize) {
- // trainable_parameters: after freezing, this is just the head.
- let mut opt = Adam::new(model.trainable_parameters(), 1e-2);
+ let mut opt = Adam::new(model.trainable_parameters(), LR);
for _ in 0..steps {
let loss = cross_entropy(&model.forward(inputs), labels);
opt.zero_grad();
@@ -66,32 +64,30 @@ fn fit(model: &impl Module, inputs: &Tensor, labels: &[usize], steps: usize) {
}
fn main() -> fastnn::Result<()> {
- manual_seed(11);
+ manual_seed(SEED);
let checkpoint = std::env::temp_dir().join("finetune_backbone.fdl");
- // ── Pretrain on task A: which quadrant is the point in? ──────────────────
let quadrant = |x: f32, y: f32| (x > 0.0) as usize + 2 * ((y > 0.0) as usize);
- let (inputs, labels) = dataset(512, quadrant);
+ let (inputs, labels) = dataset(SAMPLES, quadrant);
let pretrained = Net::new(4);
- fit(&pretrained, &inputs, &labels, 300);
+ fit(&pretrained, &inputs, &labels, STEPS);
println!(
"task A (quadrants): {:.1}%",
100.0 * accuracy(&pretrained, &inputs, &labels)
);
save(&pretrained, &checkpoint)?;
- // ── Fine-tune on task B: inside or outside the unit circle? ──────────────
let circle = |x: f32, y: f32| (x * x + y * y < 1.0) as usize;
- let (inputs, labels) = dataset(512, circle);
+ let (inputs, labels) = dataset(SAMPLES, circle);
let mut model = Net::new(4);
load(&model, &checkpoint)?;
model.backbone.freeze();
- model.head = Linear::new(16, 2); // fresh head for the new label space
+ model.head = Linear::new(HIDDEN, 2);
let before = accuracy(&model, &inputs, &labels);
- fit(&model, &inputs, &labels, 300);
+ fit(&model, &inputs, &labels, STEPS);
println!("task B before head: {:.1}%", 100.0 * before);
println!(
"task B after head: {:.1}%",
diff --git a/examples/gan.rs b/examples/gan.rs
index 0d16c6a..2fb5863 100644
--- a/examples/gan.rs
+++ b/examples/gan.rs
@@ -1,45 +1,43 @@
-//! A minimal GAN: two networks trained against each other.
+//! Minimal GAN that learns to sample from N(3, 0.5).
//!
//! cargo run --example gan --release
-//!
-//! The generator turns noise into numbers; the discriminator guesses which
-//! numbers came from the real distribution, N(3, 0.5). Each is the other's
-//! loss function. The detach on generated samples during the discriminator
-//! step is the whole trick: it stops the discriminator's loss from reaching
-//! back into the generator's weights.
use fastnn::prelude::*;
const NOISE: usize = 8;
const BATCH: usize = 64;
+const HIDDEN: usize = 32;
+const LR: f32 = 1e-3;
+const STEPS: usize = 2000;
+const REPORT_EVERY: usize = 400;
+const SEED: u64 = 21;
fn main() {
- manual_seed(21);
+ manual_seed(SEED);
let (target_mean, target_std) = (3.0, 0.5);
let generator = Sequential::new()
- .add(Linear::new(NOISE, 32))
+ .add(Linear::new(NOISE, HIDDEN))
.add(Tanh)
- .add(Linear::new(32, 1));
+ .add(Linear::new(HIDDEN, 1));
let discriminator = Sequential::new()
- .add(Linear::new(1, 32))
+ .add(Linear::new(1, HIDDEN))
.add(LeakyReLU::new(0.2))
- .add(Linear::new(32, 1));
+ .add(Linear::new(HIDDEN, 1));
- let mut g_opt = Adam::new(generator.parameters(), 1e-3);
- let mut d_opt = Adam::new(discriminator.parameters(), 1e-3);
+ let mut g_opt = Adam::new(generator.parameters(), LR);
+ let mut d_opt = Adam::new(discriminator.parameters(), LR);
let real_labels = Tensor::from_vec(vec![1.0; BATCH], &[BATCH, 1]);
let fake_labels = Tensor::from_vec(vec![0.0; BATCH], &[BATCH, 1]);
- for step in 0..2000 {
+ for step in 0..STEPS {
let real = Tensor::randn(&[BATCH, 1])
.mul_scalar(target_std)
.add_scalar(target_mean);
let fake = generator.forward(&Tensor::randn(&[BATCH, NOISE]));
- // Discriminator: real → 1, fake → 0. The generator must not receive
- // this gradient, so its samples enter detached.
+ // Detached so the discriminator loss doesn't reach the generator.
let d_loss = bce_with_logits(&discriminator.forward(&real), &real_labels).add(
&bce_with_logits(&discriminator.forward(&fake.detach()), &fake_labels),
);
@@ -47,13 +45,12 @@ fn main() {
d_loss.backward();
d_opt.step();
- // Generator: make the discriminator call the same fakes real.
let g_loss = bce_with_logits(&discriminator.forward(&fake), &real_labels);
g_opt.zero_grad();
g_loss.backward();
g_opt.step();
- if (step + 1) % 400 == 0 {
+ if (step + 1) % REPORT_EVERY == 0 {
let sample = no_grad(|| generator.forward(&Tensor::randn(&[256, NOISE]))).to_vec();
let mean = sample.iter().sum::() / sample.len() as f32;
let var = sample.iter().map(|v| (v - mean).powi(2)).sum::() / sample.len() as f32;
diff --git a/examples/mnist_cnn.rs b/examples/mnist_cnn.rs
index aa682e0..0447c80 100644
--- a/examples/mnist_cnn.rs
+++ b/examples/mnist_cnn.rs
@@ -1,21 +1,19 @@
-//! MNIST with a small convolutional network.
+//! MNIST with two conv blocks and a linear head, around 99% test accuracy.
//!
-//! cargo run --example mnist_cnn --release
-//! cargo run --example mnist_cnn --release --features cuda
-//!
-//! Two conv blocks (convolution, batch norm, ReLU, max pool) then a linear head.
-//! Slower per step than the MLP but reaches ~99% test accuracy, because the
-//! convolutions reuse the same filters at every position instead of learning
-//! each pixel independently.
+//! cargo run --example mnist_cnn --release [--features cuda]
use fastnn::data::{Mnist, Split};
use fastnn::prelude::*;
const EPOCHS: usize = 2;
const BATCH: usize = 64;
+const TEST_BATCH: usize = 500;
+const LR: f32 = 1e-3;
+const SEED: u64 = 42;
+const REPORT_EVERY: usize = 100;
fn main() -> fastnn::Result<()> {
- manual_seed(42);
+ manual_seed(SEED);
let train = Mnist::load(Split::Train)?;
let test = Mnist::load(Split::Test)?;
@@ -36,13 +34,13 @@ fn main() -> fastnn::Result<()> {
.add(Linear::new(32 * 7 * 7, 10));
model.to_device(device);
- let mut opt = Adam::new(model.parameters(), 1e-3);
- let schedule = CosineAnnealing::new(1e-3, EPOCHS * (train.len() / BATCH));
+ let mut opt = Adam::new(model.parameters(), LR);
+ let schedule = CosineAnnealing::new(LR, EPOCHS * (train.len() / BATCH));
let train_loader = DataLoader::new(&train, BATCH)
.shuffle(true)
.drop_last(true)
.to_device(device);
- let test_loader = DataLoader::new(&test, 500).to_device(device);
+ let test_loader = DataLoader::new(&test, TEST_BATCH).to_device(device);
println!("model: {} parameters on {device}\n", model.num_parameters());
@@ -67,7 +65,7 @@ fn main() -> fastnn::Result<()> {
correct += count_correct(&logits, &labels);
seen += batch.len();
- if (index + 1) % 100 == 0 {
+ if (index + 1) % REPORT_EVERY == 0 {
println!(
" epoch {epoch} batch {:4}/{} loss {:.4} acc {:.2}% lr {:.2e}",
index + 1,
@@ -79,8 +77,7 @@ fn main() -> fastnn::Result<()> {
}
}
- // eval() switches batch norm to its running statistics, so a single
- // image scores the same as it would inside a full batch.
+ // eval() makes batch norm use its running stats.
model.eval();
println!(
"epoch {epoch} done test acc {:.2}%\n",
diff --git a/examples/mnist_mlp.rs b/examples/mnist_mlp.rs
index 547527d..340ff67 100644
--- a/examples/mnist_mlp.rs
+++ b/examples/mnist_mlp.rs
@@ -1,18 +1,21 @@
-//! MNIST with a two-layer MLP. Downloads the dataset on first run.
+//! MNIST with a two-layer MLP, around 97% test accuracy. Downloads the dataset
+//! on first run.
//!
-//! cargo run --example mnist_mlp --release
-//! cargo run --example mnist_mlp --release --features cuda
-//!
-//! Reaches roughly 97% test accuracy in three epochs.
+//! cargo run --example mnist_mlp --release [--features cuda]
use fastnn::data::{Mnist, Split};
use fastnn::prelude::*;
const EPOCHS: usize = 3;
const BATCH: usize = 128;
+const TEST_BATCH: usize = 1000;
+const HIDDEN: usize = 128;
+const LR: f32 = 1e-3;
+const SEED: u64 = 42;
+const REPORT_EVERY: usize = 100;
fn main() -> fastnn::Result<()> {
- manual_seed(42);
+ manual_seed(SEED);
let train = Mnist::load(Split::Train)?;
let test = Mnist::load(Split::Test)?;
@@ -25,16 +28,16 @@ fn main() -> fastnn::Result<()> {
let device = Device::best();
let model = Sequential::new()
.add(Flatten::new())
- .add(Linear::new(784, 128))
+ .add(Linear::new(784, HIDDEN))
.add(ReLU)
- .add(Linear::new(128, 10));
+ .add(Linear::new(HIDDEN, 10));
model.to_device(device);
- let mut opt = Adam::new(model.parameters(), 1e-3);
+ let mut opt = Adam::new(model.parameters(), LR);
let train_loader = DataLoader::new(&train, BATCH)
.shuffle(true)
.to_device(device);
- let test_loader = DataLoader::new(&test, 1000).to_device(device);
+ let test_loader = DataLoader::new(&test, TEST_BATCH).to_device(device);
println!("model: {} parameters on {device}\n", model.num_parameters());
@@ -55,7 +58,7 @@ fn main() -> fastnn::Result<()> {
correct += count_correct(&logits, &labels);
seen += batch.len();
- if (index + 1) % 100 == 0 {
+ if (index + 1) % REPORT_EVERY == 0 {
println!(
" epoch {epoch} batch {:4}/{} loss {:.4} acc {:.2}%",
index + 1,
@@ -79,7 +82,6 @@ fn main() -> fastnn::Result<()> {
Ok(())
}
-/// Fraction of the loader's items the model classifies correctly.
fn accuracy(model: &dyn Module, loader: &DataLoader) -> f32 {
let (mut correct, mut seen) = (0usize, 0usize);
for batch in loader.iter() {
diff --git a/examples/multi_task.rs b/examples/multi_task.rs
index 38b1c40..a81e57b 100644
--- a/examples/multi_task.rs
+++ b/examples/multi_task.rs
@@ -1,16 +1,18 @@
-//! Multi-task learning: one shared trunk, two heads, one combined loss.
+//! Multi-task learning: a shared trunk with two heads and a summed loss.
+//! One head predicts the sign of x, the other inside/outside the unit circle.
//!
//! cargo run --example multi_task --release
-//!
-//! Both tasks read the same 2-D point: one head classifies its sign along x,
-//! the other whether it falls inside the unit circle. The trunk must learn
-//! features that serve both at once — the sum of the two losses is all it
-//! takes, because gradients from each head simply accumulate in the shared
-//! parameters.
use fastnn::nn::module::scoped;
use fastnn::prelude::*;
+const SEED: u64 = 51;
+const SAMPLES: usize = 512;
+const HIDDEN: usize = 24;
+const LR: f32 = 1e-2;
+const STEPS: usize = 400;
+const REPORT_EVERY: usize = 100;
+
struct MultiTask {
trunk: Sequential,
sign_head: Linear,
@@ -21,12 +23,12 @@ impl MultiTask {
fn new() -> MultiTask {
MultiTask {
trunk: Sequential::new()
- .add(Linear::new(2, 24))
+ .add(Linear::new(2, HIDDEN))
.add(Tanh)
- .add(Linear::new(24, 24))
+ .add(Linear::new(HIDDEN, HIDDEN))
.add(Tanh),
- sign_head: Linear::new(24, 2),
- circle_head: Linear::new(24, 2),
+ sign_head: Linear::new(HIDDEN, 2),
+ circle_head: Linear::new(HIDDEN, 2),
}
}
@@ -62,19 +64,19 @@ fn accuracy(logits: &Tensor, labels: &[usize]) -> f32 {
}
fn main() {
- manual_seed(51);
- let coords = fastnn::rng::uniform(512 * 2, -1.5, 1.5);
+ manual_seed(SEED);
+ let coords = fastnn::rng::uniform(SAMPLES * 2, -1.5, 1.5);
let sign_labels: Vec = coords.chunks(2).map(|p| (p[0] > 0.0) as usize).collect();
let circle_labels: Vec = coords
.chunks(2)
.map(|p| (p[0] * p[0] + p[1] * p[1] < 1.0) as usize)
.collect();
- let inputs = Tensor::from_vec(coords, &[512, 2]);
+ let inputs = Tensor::from_vec(coords, &[SAMPLES, 2]);
let model = MultiTask::new();
- let mut opt = Adam::new(model.parameters(), 1e-2);
+ let mut opt = Adam::new(model.parameters(), LR);
- for step in 0..400 {
+ for step in 0..STEPS {
let (sign_logits, circle_logits) = model.heads(&inputs);
let loss = cross_entropy(&sign_logits, &sign_labels)
.add(&cross_entropy(&circle_logits, &circle_labels));
@@ -83,7 +85,7 @@ fn main() {
loss.backward();
opt.step();
- if (step + 1) % 100 == 0 {
+ if (step + 1) % REPORT_EVERY == 0 {
println!(
"step {:3} loss {:.4} sign {:.1}% circle {:.1}%",
step + 1,
diff --git a/examples/reinforce.rs b/examples/reinforce.rs
index 5a85fd3..438b0ff 100644
--- a/examples/reinforce.rs
+++ b/examples/reinforce.rs
@@ -1,40 +1,44 @@
-//! Policy-gradient reinforcement learning on a multi-armed bandit.
+//! REINFORCE on a five-armed bandit. The policy is a softmax over per-arm
+//! preferences, trained on advantage over a running baseline.
//!
//! cargo run --example reinforce --release
-//!
-//! REINFORCE in its smallest complete form. The policy is a softmax over five
-//! arm preferences; each pull nudges the log-probability of the chosen arm by
-//! the reward's advantage over a running baseline. `cross_entropy` already *is*
-//! negative log-likelihood, so scaling it by the advantage is the entire loss.
use fastnn::nn::Param;
use fastnn::prelude::*;
+const SEED: u64 = 41;
+const STEPS: usize = 2000;
+const LR: f32 = 5e-2;
+const REWARD_NOISE: f32 = 0.1;
+const BASELINE_RATE: f32 = 0.01;
+const REPORT_EVERY: usize = 400;
+
fn main() {
- manual_seed(41);
- // Expected payout per arm; arm 3 is best, but only by observation.
+ manual_seed(SEED);
+ // Mean payout per arm; arm 3 is the best.
let arms = [0.2f32, 0.5, 0.4, 0.9, 0.6];
let preferences = Param::new(Tensor::zeros(&[1, arms.len()]));
- let mut opt = Adam::new(vec![preferences.clone()], 5e-2);
+ let mut opt = Adam::new(vec![preferences.clone()], LR);
let sampler = Sampler::new();
let mut baseline = 0.0f32;
- for step in 0..2000 {
+ for step in 0..STEPS {
let logits = preferences.tensor();
let arm = sampler.sample(&logits.detach());
- let reward = arms[arm] + 0.1 * fastnn::rng::normal(1)[0];
+ let reward = arms[arm] + REWARD_NOISE * fastnn::rng::normal(1)[0];
- // Better than the baseline → make this arm likelier; worse → less.
+ // cross_entropy is the negative log-prob, so scaling by advantage
+ // gives the policy gradient loss.
let advantage = reward - baseline;
- baseline += 0.01 * (reward - baseline);
+ baseline += BASELINE_RATE * (reward - baseline);
let loss = cross_entropy(&logits, &[arm]).mul_scalar(advantage);
opt.zero_grad();
loss.backward();
opt.step();
- if (step + 1) % 400 == 0 {
+ if (step + 1) % REPORT_EVERY == 0 {
let probs = no_grad(|| preferences.tensor().softmax()).to_vec();
let shown: Vec = probs.iter().map(|p| format!("{p:.2}")).collect();
println!(
diff --git a/examples/simple_mlp.rs b/examples/simple_mlp.rs
index 1c12ca0..0fea773 100644
--- a/examples/simple_mlp.rs
+++ b/examples/simple_mlp.rs
@@ -1,43 +1,44 @@
-//! The smallest complete training loop: an MLP learning XOR.
+//! The smallest training loop: an MLP learning XOR.
//!
//! cargo run --example simple_mlp --release
-//!
-//! XOR is not linearly separable, so a single linear layer cannot fit it at all.
-//! Two layers with a non-linearity between them can, which makes this the
-//! shortest honest test that a network is really learning.
use fastnn::prelude::*;
+const SEED: u64 = 42;
+const HIDDEN: usize = 8;
+const LR: f32 = 0.05;
+const STEPS: usize = 400;
+const REPORT_EVERY: usize = 100;
+
fn main() {
- manual_seed(42);
+ manual_seed(SEED);
let inputs = Tensor::from_vec(vec![0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0], &[4, 2]);
let targets = [0usize, 1, 1, 0];
let model = Sequential::new()
- .add(Linear::new(2, 8))
+ .add(Linear::new(2, HIDDEN))
.add(Tanh)
- .add(Linear::new(8, 2));
+ .add(Linear::new(HIDDEN, 2));
- let mut opt = Adam::new(model.parameters(), 0.05);
+ let mut opt = Adam::new(model.parameters(), LR);
println!(
"training a {}-parameter MLP on XOR\n",
model.num_parameters()
);
- for step in 1..=400 {
+ for step in 1..=STEPS {
let loss = cross_entropy(&model.forward(&inputs), &targets);
opt.zero_grad();
loss.backward();
opt.step();
- if step % 100 == 0 {
+ if step % REPORT_EVERY == 0 {
println!(" step {step:3} loss {:.4}", loss.item());
}
}
- // Inference builds no graph, so skip it.
let predictions = no_grad(|| model.forward(&inputs)).argmax(1);
let values = inputs.to_vec();
diff --git a/examples/vae.rs b/examples/vae.rs
index b3d543f..e0a7282 100644
--- a/examples/vae.rs
+++ b/examples/vae.rs
@@ -1,18 +1,19 @@
-//! A variational autoencoder on 2-D data.
+//! VAE on noisy points around the unit circle.
//!
//! cargo run --example vae --release
-//!
-//! The data lives on a noisy circle. The encoder outputs a *distribution*
-//! (mean and log-variance) rather than a point; the reparameterisation trick —
-//! `z = mu + eps · exp(logvar/2)` with `eps` drawn outside the graph — is what
-//! lets the sampling step stay differentiable. The KL term pulls the latent
-//! space toward N(0, I), which is why sampling from the prior afterwards
-//! produces data-like points.
use fastnn::nn::module::scoped;
use fastnn::prelude::*;
const LATENT: usize = 2;
+const HIDDEN: usize = 32;
+const SEED: u64 = 31;
+const BATCH: usize = 128;
+const LR: f32 = 1e-3;
+const STEPS: usize = 3000;
+const REPORT_EVERY: usize = 600;
+const KL_WEIGHT: f32 = 0.1;
+const PRIOR_SAMPLES: usize = 256;
struct Vae {
encoder: Sequential,
@@ -24,22 +25,23 @@ struct Vae {
impl Vae {
fn new() -> Vae {
Vae {
- encoder: Sequential::new().add(Linear::new(2, 32)).add(Tanh),
- to_mean: Linear::new(32, LATENT),
- to_logvar: Linear::new(32, LATENT),
+ encoder: Sequential::new().add(Linear::new(2, HIDDEN)).add(Tanh),
+ to_mean: Linear::new(HIDDEN, LATENT),
+ to_logvar: Linear::new(HIDDEN, LATENT),
decoder: Sequential::new()
- .add(Linear::new(LATENT, 32))
+ .add(Linear::new(LATENT, HIDDEN))
.add(Tanh)
- .add(Linear::new(32, 2)),
+ .add(Linear::new(HIDDEN, 2)),
}
}
- /// Reconstruction plus the two pieces the loss needs.
+ /// Returns the reconstruction, mean and log-variance.
fn forward_full(&self, x: &Tensor) -> (Tensor, Tensor, Tensor) {
let hidden = self.encoder.forward(x);
let mean = self.to_mean.forward(&hidden);
let logvar = self.to_logvar.forward(&hidden);
+ // eps is sampled outside the graph so z stays differentiable.
let eps = Tensor::randn(&[x.dim(0), LATENT]);
let z = mean.add(&logvar.mul_scalar(0.5).exp().mul(&eps));
(self.decoder.forward(&z), mean, logvar)
@@ -59,7 +61,6 @@ impl Module for Vae {
}
}
-/// Points on the unit circle plus noise.
fn circle_batch(n: usize) -> Tensor {
let angles = fastnn::rng::uniform(n, 0.0, std::f32::consts::TAU);
let noise = fastnn::rng::normal(n * 2);
@@ -72,16 +73,16 @@ fn circle_batch(n: usize) -> Tensor {
}
fn main() {
- manual_seed(31);
+ manual_seed(SEED);
let model = Vae::new();
- let mut opt = Adam::new(model.parameters(), 1e-3);
+ let mut opt = Adam::new(model.parameters(), LR);
- for step in 0..3000 {
- let batch = circle_batch(128);
+ for step in 0..STEPS {
+ let batch = circle_batch(BATCH);
let (reconstructed, mean, logvar) = model.forward_full(&batch);
let reconstruction = mse(&reconstructed, &batch);
- // KL(N(mu, sigma) ‖ N(0, 1)), summed over latent dims, averaged over the batch.
+ // KL to N(0, 1), summed over latent dims and averaged over the batch.
let kl = mean
.square()
.add(&logvar.exp())
@@ -90,13 +91,13 @@ fn main() {
.mul_scalar(0.5)
.sum()
.div_scalar(batch.dim(0) as f32);
- let loss = reconstruction.add(&kl.mul_scalar(0.1));
+ let loss = reconstruction.add(&kl.mul_scalar(KL_WEIGHT));
opt.zero_grad();
loss.backward();
opt.step();
- if (step + 1) % 600 == 0 {
+ if (step + 1) % REPORT_EVERY == 0 {
println!(
"step {:4} reconstruction {:.4} kl {:.3}",
step + 1,
@@ -106,13 +107,19 @@ fn main() {
}
}
- // Decode from the prior: no encoder, no data — radius ≈ 1 means the latent
- // space really did learn the circle.
- let generated = no_grad(|| model.decoder.forward(&Tensor::randn(&[256, LATENT]))).to_vec();
+ // Decode from the prior; a radius near 1 means it learned the circle.
+ let generated = no_grad(|| {
+ model
+ .decoder
+ .forward(&Tensor::randn(&[PRIOR_SAMPLES, LATENT]))
+ })
+ .to_vec();
let mean_radius: f32 = generated
.chunks(2)
.map(|p| (p[0] * p[0] + p[1] * p[1]).sqrt())
.sum::()
- / 256.0;
- println!("\nmean radius of 256 prior samples: {mean_radius:.3} (data lives at 1.0)");
+ / PRIOR_SAMPLES as f32;
+ println!(
+ "\nmean radius of {PRIOR_SAMPLES} prior samples: {mean_radius:.3} (data lives at 1.0)"
+ );
}
diff --git a/src/autograd/anomaly.rs b/src/autograd/anomaly.rs
index 0a90b58..3569f03 100644
--- a/src/autograd/anomaly.rs
+++ b/src/autograd/anomaly.rs
@@ -1,26 +1,6 @@
-//! Catching a NaN at the op that produced it.
-//!
-//! A single non-finite gradient spreads: the optimizer writes it into a
-//! parameter, the next forward pass turns every activation downstream into NaN,
-//! and the loss only *prints* as NaN some steps later — by which point the
-//! checkpoint on disk is poisoned too and the op actually at fault is long gone.
-//!
-//! Wrapping a step in [`detect_anomaly`] makes the reverse pass check each rule's
-//! output and fail immediately, naming the rule.
-//!
-//! ```should_panic
-//! # use fastnn::prelude::*;
-//! # use fastnn::autograd::detect_anomaly;
-//! let x = Param::new(Tensor::zeros(&[2, 2]));
-//!
-//! detect_anomaly(|| {
-//! // d/dx log(x) = 1/x, which is infinite at x = 0
-//! x.tensor().log().sum().backward(); // panics naming "Log"
-//! });
-//! ```
-//!
-//! It is a debugging tool, not something to leave on: the check reads every
-//! gradient, which on a GPU means downloading each one to the host.
+//! `detect_anomaly`: panic at the backward rule that first produces a NaN or inf,
+//! instead of letting it reach the optimizer. It downloads every gradient on a
+//! GPU, so leave it off outside debugging.
use std::cell::Cell;
@@ -35,15 +15,13 @@ pub fn is_detecting() -> bool {
ENABLED.with(|flag| flag.get())
}
-/// Run `f` with gradient checking on, then restore the previous setting.
+/// Run `f` with gradient checking on.
pub fn detect_anomaly(f: impl FnOnce() -> T) -> T {
let _guard = Enabled::new();
f()
}
-/// Panic if `grads` holds a non-finite value, naming `op` and which input it was.
-///
-/// Called by the engine after each rule while detection is on.
+/// Panic if `grads` holds a non-finite value, naming the op and input.
pub(crate) fn check(op: &str, grads: &[Tensor]) {
for (index, grad) in grads.iter().enumerate() {
if let Some(report) = first_bad(grad) {
@@ -52,7 +30,6 @@ pub(crate) fn check(op: &str, grads: &[Tensor]) {
}
}
-/// Describe the first non-finite value in `t`, if there is one.
fn first_bad(t: &Tensor) -> Option {
let data = t.to_vec();
let (slot, value) = data.iter().enumerate().find(|(_, v)| !v.is_finite())?;
@@ -60,8 +37,7 @@ fn first_bad(t: &Tensor) -> Option {
Some(format!("{kind} at element {slot} of {:?}", t.shape()))
}
-/// Restores the previous setting on drop, so a panic inside `detect_anomaly`
-/// cannot leave checking permanently on.
+/// Guard, so a panic inside `detect_anomaly` cannot leave checking on.
struct Enabled(bool);
impl Enabled {
@@ -77,10 +53,7 @@ impl Drop for Enabled {
}
impl Tensor {
- /// Whether every element is finite — no NaN, no infinity.
- ///
- /// Reads the whole tensor, and downloads it first if it is on a GPU. Use it
- /// at checkpoints and guard rails, not inside a training loop.
+ /// Whether every element is finite. Reads the whole tensor, so keep it out of hot loops.
pub fn is_finite(&self) -> bool {
self.to_vec().iter().all(|v| v.is_finite())
}
diff --git a/src/autograd/engine.rs b/src/autograd/engine.rs
index 89e1e45..fc9d7d5 100644
--- a/src/autograd/engine.rs
+++ b/src/autograd/engine.rs
@@ -1,4 +1,4 @@
-//! The reverse pass: walk the graph from a scalar loss back to the leaves.
+//! The reverse pass from a scalar loss back to the leaves.
use std::collections::{HashMap, HashSet};
@@ -6,14 +6,11 @@ use crate::tensor::Tensor;
use super::node::{Node, NodeId};
-/// Seed `loss` with a gradient of 1 and propagate it to every reachable leaf.
+/// Propagate a gradient of 1 from `loss` to every reachable leaf.
///
-/// Gradients land in each parameter's slot, additively — call
-/// [`Optimizer::zero_grad`](crate::optim::Optimizer::zero_grad) between steps, or
-/// skip it to accumulate across micro-batches.
-///
-/// Panics if `loss` is not a scalar. Does nothing if `loss` carries no graph
-/// (a constant, or built inside [`no_grad`](super::no_grad)).
+/// Gradients add into the parameter slots, so call
+/// [`Optimizer::zero_grad`](crate::optim::Optimizer::zero_grad) between steps.
+/// Panics if `loss` is not a scalar, does nothing if it has no graph.
pub fn backward(loss: &Tensor) {
assert_eq!(
loss.numel(),
@@ -23,8 +20,7 @@ pub fn backward(loss: &Tensor) {
);
let Some(root) = loss.grad_fn() else { return };
- // Seed on the loss's own device: every rule derives its output's device from
- // the incoming gradient, so seeding on the wrong one poisons the whole pass.
+ // rules take their device from the incoming gradient, so seed on the loss's
let mut grads: HashMap = HashMap::new();
grads.insert(root.id(), Tensor::ones(&[1]).to(loss.device()));
@@ -41,8 +37,7 @@ pub fn backward(loss: &Tensor) {
Node::Op(op) => op,
};
- // Rules compute with ordinary tensor ops. Suppressing tracking here means
- // a rule that forgets to detach a saved value still cannot grow a graph.
+ // no_grad so a rule that forgets to detach a saved value cannot grow a graph
let input_grads = super::no_grad(|| op.rule.backward(&grad));
if super::anomaly::is_detecting() {
@@ -69,10 +64,8 @@ pub fn backward(loss: &Tensor) {
}
}
-/// Nodes ordered so that every consumer comes before the values it consumed.
-///
-/// Depth-first post-order reversed. The traversal is iterative because a
-/// transformer's graph is deeper than a comfortable recursion limit.
+/// Nodes with every consumer before the values it consumed. Iterative, since
+/// deep graphs would overflow the stack.
fn reverse_topological(root: &Node) -> Vec {
let mut order = Vec::new();
let mut visited = HashSet::new();
diff --git a/src/autograd/mod.rs b/src/autograd/mod.rs
index 56e38a0..401d553 100644
--- a/src/autograd/mod.rs
+++ b/src/autograd/mod.rs
@@ -1,26 +1,5 @@
-//! Reverse-mode automatic differentiation.
-//!
-//! There is no tape and no global state. Every tensor produced by a differentiable
-//! op carries a [`Node`] naming the rule that made it and the tensors it consumed,
-//! so the graph *is* the tensors. It is freed when the loss is dropped.
-//!
-//! ```text
-//! w (Leaf) ─┐
-//! ├─► matmul ──► add ──► relu ──► … ──► loss
-//! x (None) ─┘ ▲ ▲
-//! │ └─ Node::Op { rule: AddBackward, inputs: [.., b] }
-//! └────────── Node::Op { rule: MatmulBackward, inputs: [x, w] }
-//! ```
-//!
-//! `loss.backward()` walks that structure in reverse and deposits gradients in
-//! the [`GradSlot`] of each leaf, which is the same slot the optimizer reads.
-//!
-//! ## Adding an op
-//!
-//! 1. Write the forward in `tensor/ops/`.
-//! 2. Write a `Backward` rule in [`ops`], saving only what the derivative needs.
-//! 3. Attach it with `Tensor::with_grad`.
-//! 4. Add a finite-difference check in `tests/gradcheck.rs`.
+//! Reverse-mode autodiff. Each differentiable tensor carries a [`Node`] with its
+//! backward rule and inputs, so the graph is just the tensors and there is no tape.
pub mod anomaly;
pub mod engine;
diff --git a/src/autograd/mode.rs b/src/autograd/mode.rs
index ec8506f..be1697b 100644
--- a/src/autograd/mode.rs
+++ b/src/autograd/mode.rs
@@ -1,13 +1,8 @@
-//! Whether ops attach graph nodes to their outputs.
-//!
-//! Tracking is on by default. Wrap inference in [`no_grad`] to skip building a
-//! graph you will never differentiate — it saves the memory the saved forward
-//! values would otherwise hold.
+//! The `no_grad` switch: whether ops attach graph nodes to their outputs.
use std::cell::Cell;
thread_local! {
- /// Nesting depth of active `no_grad` scopes. Zero means "record".
static SUPPRESSED: Cell = const { Cell::new(0) };
}
@@ -16,7 +11,7 @@ pub fn is_enabled() -> bool {
SUPPRESSED.with(|s| s.get() == 0)
}
-/// Run `f` without building a graph, then restore the previous mode.
+/// Run `f` without building a graph.
///
/// ```
/// # use fastnn::prelude::*;
@@ -30,8 +25,7 @@ pub fn no_grad(f: impl FnOnce() -> T) -> T {
f()
}
-/// Restores the previous mode on drop, so a panic inside `no_grad` cannot leave
-/// tracking permanently off.
+/// Guard, so a panic inside `no_grad` cannot leave tracking off.
struct Suppress;
impl Suppress {
diff --git a/src/autograd/node.rs b/src/autograd/node.rs
index 547e7cc..35e6b3f 100644
--- a/src/autograd/node.rs
+++ b/src/autograd/node.rs
@@ -4,33 +4,29 @@ use std::sync::{Arc, Mutex};
use crate::tensor::Tensor;
-/// How one operation turns the gradient of its output into gradients of its inputs.
+/// Turns the gradient of an op's output into gradients of its inputs.
///
-/// A rule owns whatever the forward pass needs to save (operands, masks, pooling
-/// indices). Save *detached* tensors — a saved value that still carries its own
-/// node would pin the graph that produced it.
+/// Save detached tensors in a rule, or the saved value pins the graph behind it.
pub trait Backward: Send + Sync {
- /// Gradients for each input, in the same order as [`Op::inputs`].
+ /// One gradient per input, in the order of [`Op::inputs`].
fn backward(&self, grad: &Tensor) -> Vec;
- /// Shown in panics and graph dumps.
+ /// Used in panic messages.
fn name(&self) -> &'static str;
}
-/// An operation that produced a tensor, plus the tensors it consumed.
+/// The rule that produced a tensor and the tensors it consumed.
pub struct Op {
pub rule: Box,
pub inputs: Vec,
}
-/// A tensor's link into the graph. `None` on a tensor means "constant data" —
-/// gradient flow stops there.
+/// A tensor's link into the graph. Tensors without one are constants.
#[derive(Clone)]
pub enum Node {
- /// A learnable value. Gradient reaching it lands in the shared slot, which
- /// is the same slot the optimizer reads.
+ /// A parameter. Its gradient lands in the slot the optimizer reads.
Leaf(GradSlot),
- /// An intermediate value. Gradient reaching it is pushed further back.
+ /// An intermediate value.
Op(Arc),
}
@@ -42,10 +38,8 @@ impl Node {
}))
}
- /// Stable identity, used to accumulate gradients per node during backward.
- ///
- /// Two handles to the same parameter share a `GradSlot` and therefore an id,
- /// which is exactly what makes weight tying work.
+ /// Identity for accumulating gradients. Handles to the same parameter share
+ /// a slot and so an id, which is what makes weight tying work.
pub fn id(&self) -> NodeId {
match self {
Node::Leaf(slot) => NodeId(Arc::as_ptr(&slot.0) as *const () as usize),
@@ -57,7 +51,7 @@ impl Node {
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub struct NodeId(usize);
-/// The gradient of one learnable tensor, shared between the model and the optimizer.
+/// A parameter's gradient, shared between the model and the optimizer.
#[derive(Clone, Default)]
pub struct GradSlot(Arc>>);
@@ -66,11 +60,8 @@ impl GradSlot {
GradSlot::default()
}
- /// Add `grad` to whatever is already there.
- ///
- /// Accumulating rather than overwriting is what lets a parameter be used
- /// twice in one forward pass, and what makes gradient accumulation across
- /// micro-batches work without extra bookkeeping.
+ /// Add to the stored gradient. Adding (not overwriting) covers parameters
+ /// used twice and accumulation across micro-batches.
pub fn accumulate(&self, grad: Tensor) {
let mut slot = self.0.lock().unwrap();
*slot = Some(match slot.take() {
diff --git a/src/autograd/ops/activation.rs b/src/autograd/ops/activation.rs
index 8c136ff..b0d8c96 100644
--- a/src/autograd/ops/activation.rs
+++ b/src/autograd/ops/activation.rs
@@ -1,8 +1,5 @@
-//! Derivatives of activations and the softmax family.
-//!
-//! Where `kernels.cu` has a fused backward kernel we use it; the host formula
-//! below it is both the CPU path and the reference the GPU kernel must match
-//! (see `tests/cuda_parity.rs`).
+//! Derivatives of activations and softmax. Fused GPU kernels are used where they
+//! exist, and the host formulas are the CPU path and the parity reference.
use crate::autograd::Backward;
use crate::cuda::{ffi, kernels};
@@ -11,7 +8,7 @@ use crate::tensor::Tensor;
use super::unary::elementwise;
-/// Run a fused `(grad, saved) → grad_input` kernel when both live on the GPU.
+/// Run a fused backward kernel if both tensors are on the GPU.
fn fused(grad: &Tensor, saved: &Tensor, kernel: ffi::BinaryKernel) -> Option {
let (Storage::Cuda(g), Storage::Cuda(s)) = (grad.storage(), saved.storage()) else {
return None;
@@ -40,7 +37,7 @@ impl Backward for ReluBackward {
}
}
-/// `d σ(x) = g · σ(1 - σ)` — in terms of the forward output.
+/// `d σ(x) = g · σ(1 - σ)`, in terms of the forward output.
pub struct SigmoidBackward {
pub output: Tensor,
}
@@ -75,7 +72,7 @@ impl Backward for TanhBackward {
}
}
-/// `d gelu(x) = g · (Φ(x) + x·φ(x))`, the exact form matching the forward.
+/// `d gelu(x) = g · (Φ(x) + x·φ(x))`.
pub struct GeluBackward {
pub input: Tensor,
}
@@ -129,8 +126,7 @@ impl Backward for LeakyReluBackward {
}
}
-/// `d softmax = s · (g - Σ(g·s))` per row — the Jacobian collapses to this
-/// because every output in a row depends on every input in that row.
+/// `d softmax = s · (g - Σ(g·s))` per row.
pub struct SoftmaxBackward {
pub output: Tensor,
}
diff --git a/src/autograd/ops/arith.rs b/src/autograd/ops/arith.rs
index e8f7260..601e4a9 100644
--- a/src/autograd/ops/arith.rs
+++ b/src/autograd/ops/arith.rs
@@ -73,7 +73,7 @@ impl Backward for DivBackward {
}
}
-/// `d(a · k) = g·k`. With `factor = 1` this is also the rule for `+ k`.
+/// `d(a · k) = g·k`. Also used for `+ k` with `factor = 1`.
pub struct MulScalarBackward {
pub factor: f32,
}
diff --git a/src/autograd/ops/conv.rs b/src/autograd/ops/conv.rs
index 8b27405..4361b4b 100644
--- a/src/autograd/ops/conv.rs
+++ b/src/autograd/ops/conv.rs
@@ -1,10 +1,5 @@
-//! Derivatives of the `im2col` / `col2im` adjoint pair.
-//!
-//! The two operations are transposes of one linear map, so each is the
-//! other's gradient: differentiating the unfold means folding the incoming
-//! gradient, and differentiating the fold means unfolding it. Both rules
-//! delegate to the same shared index maps the forwards use — the adjoint
-//! identity holds by construction, not by keeping two loops in sync.
+//! Derivatives of `im2col` and `col2im`. Each is the other's gradient, and both
+//! use the forward's index maps so they stay exact adjoints.
use crate::autograd::Backward;
use crate::cuda::kernels;
@@ -12,10 +7,7 @@ use crate::tensor::ops::conv::{fold, unfold, Window};
use crate::tensor::storage::Storage;
use crate::tensor::Tensor;
-/// Gradient of `im2col`: fold columns back into the image, adding where
-/// windows overlapped. A pixel covered by several windows was copied into
-/// several columns, so its gradient is the sum of theirs; pixels that only
-/// ever landed in padding get nothing.
+/// Gradient of `im2col`: fold the columns back, summing overlaps.
pub struct Col2ImBackward {
pub shape: Vec,
pub window: Window,
@@ -54,10 +46,7 @@ impl Backward for Col2ImBackward {
}
}
-/// Gradient of `col2im`: unfold the incoming image gradient back into columns.
-///
-/// Each column entry contributed to exactly one pixel, so its gradient is that
-/// pixel's — which is precisely what `im2col` reads.
+/// Gradient of `col2im`: unfold the image gradient.
pub struct Im2ColBackward {
pub window: Window,
}
diff --git a/src/autograd/ops/index.rs b/src/autograd/ops/index.rs
index e77b348..5c31457 100644
--- a/src/autograd/ops/index.rs
+++ b/src/autograd/ops/index.rs
@@ -5,10 +5,7 @@ use crate::cuda::kernels;
use crate::tensor::storage::Storage;
use crate::tensor::Tensor;
-/// Scatter-add each gathered row's gradient back into its source row.
-///
-/// Rows selected more than once accumulate, which is how a token repeated in a
-/// batch contributes to its embedding once per occurrence.
+/// Scatter-add each row's gradient into its source row, so repeats accumulate.
pub struct IndexSelectBackward {
pub ids: Vec,
pub rows: usize,
diff --git a/src/autograd/ops/loss.rs b/src/autograd/ops/loss.rs
index 871b83f..523e336 100644
--- a/src/autograd/ops/loss.rs
+++ b/src/autograd/ops/loss.rs
@@ -1,13 +1,10 @@
-//! Derivative of the fused cross-entropy loss.
+//! Derivatives of the cross-entropy losses.
use crate::autograd::Backward;
use crate::tensor::Tensor;
-/// `d/dlogits mean(-log softmax(logits)[target]) = (softmax - one_hot) / batch`.
-///
-/// Fusing softmax and the negative log likelihood is what makes this one
-/// subtraction. Done separately, the same gradient would go through a log, a
-/// division, and a gather, each amplifying the rounding error of the last.
+/// `(softmax - one_hot) / batch`. Fusing softmax with the NLL makes this one
+/// subtraction instead of a log, a division and a gather.
pub struct CrossEntropyBackward {
pub softmax: Tensor,
pub targets: Vec,
@@ -34,13 +31,8 @@ impl Backward for CrossEntropyBackward {
}
}
-/// Gradient of the configurable cross-entropy.
-///
-/// The forward already computed `w·(p − q)` per row — softmax minus the
-/// (possibly smoothed) target distribution, scaled by that row's class weight,
-/// zero for ignored rows. All that remains is the reduction's scaling: the
-/// weighted-mean denominator for `Mean`, nothing for `Sum`, and a per-row
-/// upstream gradient for `None`.
+/// Configurable cross-entropy. The forward already stored `w·(p - q)` per row
+/// (zero for ignored rows), so only the reduction's scaling is left.
pub struct WeightedCrossEntropyBackward {
pub difference: Tensor,
pub normalizer: f32,
@@ -51,7 +43,6 @@ impl Backward for WeightedCrossEntropyBackward {
fn backward(&self, grad: &Tensor) -> Vec {
let difference = self.difference.to(grad.device());
let scaled = if self.per_row {
- // Upstream is one gradient per row; broadcast it across classes.
let rows = grad.numel() as i64;
difference.mul(&grad.reshape(&[rows, 1]))
} else {
diff --git a/src/autograd/ops/matmul.rs b/src/autograd/ops/matmul.rs
index e5774f7..b994265 100644
--- a/src/autograd/ops/matmul.rs
+++ b/src/autograd/ops/matmul.rs
@@ -1,8 +1,5 @@
-//! Derivatives of matrix multiplication.
-//!
-//! Each rule is two more multiplications, chosen so that neither needs a
-//! transposed copy of anything — that is the whole reason `matmul_nt` and
-//! `matmul_tn` exist as first-class ops.
+//! Derivatives of matmul. Each is two more matmuls, picked so no transposed copy
+//! is needed.
use crate::autograd::Backward;
use crate::tensor::Tensor;
@@ -60,10 +57,7 @@ impl Backward for MatmulTnBackward {
}
}
-/// Sum away any batch dimension the forward pass broadcast over.
-///
-/// A weight matrix multiplied against a batch of inputs is the usual case: its
-/// gradient comes back batched and every sample contributed to it.
+/// Sum away batch dims the forward broadcast over, e.g. a shared weight matrix.
fn unbatch(d_a: Tensor, d_b: Tensor, a: &Tensor, b: &Tensor) -> Vec {
vec![reduce_to(&d_a, a.shape()), reduce_to(&d_b, b.shape())]
}
diff --git a/src/autograd/ops/mod.rs b/src/autograd/ops/mod.rs
index a7e5624..bea5fd6 100644
--- a/src/autograd/ops/mod.rs
+++ b/src/autograd/ops/mod.rs
@@ -1,8 +1,5 @@
-//! Backward rules, mirroring the layout of [`tensor::ops`](crate::tensor::ops).
-//!
-//! A rule is a small struct holding what the derivative needs, plus one
-//! [`Backward::backward`](super::Backward::backward) method. They run with
-//! tracking suppressed, so they can use ordinary tensor ops freely.
+//! Backward rules, one file per file in [`tensor::ops`](crate::tensor::ops).
+//! They run with tracking off, so they can use ordinary tensor ops.
pub mod activation;
pub mod arith;
@@ -18,11 +15,7 @@ pub mod view;
use crate::tensor::Tensor;
-/// Undo broadcasting: sum `grad` down to `target`.
-///
-/// A forward op that broadcast a `[1, n]` bias across a batch produced a
-/// `[batch, n]` output, so its gradient arrives at `[batch, n]` and every row
-/// contributed. Summing the stretched axes is the adjoint of stretching them.
+/// Undo broadcasting by summing `grad` down to `target`.
pub(crate) fn reduce_to(grad: &Tensor, target: &[usize]) -> Tensor {
if grad.shape() == target {
return grad.clone();
diff --git a/src/autograd/ops/norm.rs b/src/autograd/ops/norm.rs
index 556fed0..cd037c7 100644
--- a/src/autograd/ops/norm.rs
+++ b/src/autograd/ops/norm.rs
@@ -1,8 +1,4 @@
-//! Derivatives of the fused normalization ops.
-//!
-//! Normalizing couples every element in a group: change one input and the mean
-//! and variance move, so every output in that group moves. The two correction
-//! terms below are exactly that coupling.
+//! Derivatives of the fused norms.
use crate::autograd::Backward;
use crate::cuda::kernels;
@@ -32,7 +28,7 @@ impl Backward for LayerNormBackward {
for r in 0..rows {
let base = r * cols;
- // Correction terms: how much the mean and the variance shift.
+ // the mean and variance both move with every input in the row
let mut mean_shift = 0.0f32;
let mut scale_shift = 0.0f32;
for c in 0..cols {
@@ -62,7 +58,7 @@ impl Backward for LayerNormBackward {
}
}
-/// LayerNorm on the GPU — one kernel produces all three gradients.
+/// LayerNorm on the GPU, one kernel for all three gradients.
pub struct LayerNormCudaBackward {
pub input: Tensor,
pub gamma: Tensor,
@@ -101,8 +97,7 @@ impl Backward for LayerNormCudaBackward {
}
}
-/// BatchNorm2d. Same structure as LayerNorm, but the group is a channel across
-/// the whole batch rather than a row across features.
+/// BatchNorm2d. Like LayerNorm but the group is a channel across the batch.
pub struct BatchNorm2dBackward {
pub normalized: Tensor,
pub gamma: Vec,
diff --git a/src/autograd/ops/pool.rs b/src/autograd/ops/pool.rs
index 2bdfc39..6f60fd6 100644
--- a/src/autograd/ops/pool.rs
+++ b/src/autograd/ops/pool.rs
@@ -1,12 +1,11 @@
-//! Derivatives of pooling — each one routes the gradient the way the forward
-//! pass routed the values.
+//! Derivatives of pooling, which route the gradient the way the forward did.
use crate::autograd::Backward;
use crate::tensor::ops::conv::Window;
use crate::tensor::ops::pool::adaptive_range;
use crate::tensor::Tensor;
-/// Only the winning input affected the output, so only it gets the gradient.
+/// Only the winning input gets the gradient.
pub struct MaxPool2dBackward {
pub shape: Vec,
pub winners: Vec,
@@ -16,7 +15,7 @@ impl Backward for MaxPool2dBackward {
fn backward(&self, grad: &Tensor) -> Vec {
let g = grad.to_vec();
let mut out = vec![0.0f32; self.shape.iter().product()];
- // `+=`, not `=`: overlapping windows can pick the same winner.
+ // overlapping windows can share a winner
for (slot, &winner) in self.winners.iter().enumerate() {
out[winner] += g[slot];
}
@@ -27,7 +26,7 @@ impl Backward for MaxPool2dBackward {
}
}
-/// Every input in the window contributed `1/count`, so each gets that share.
+/// Each input in the window gets `1/count` of the gradient.
pub struct AvgPool2dBackward {
pub shape: Vec,
pub window: Window,
@@ -58,7 +57,7 @@ impl Backward for AvgPool2dBackward {
}
}
-/// Same idea, over the input-size-dependent spans the forward pass averaged.
+/// Same, over the spans the forward averaged.
pub struct AdaptiveAvgPool2dBackward {
pub shape: Vec,
pub output: (usize, usize),
diff --git a/src/autograd/ops/reduce.rs b/src/autograd/ops/reduce.rs
index 0c5b9bb..5575871 100644
--- a/src/autograd/ops/reduce.rs
+++ b/src/autograd/ops/reduce.rs
@@ -1,13 +1,10 @@
-//! Derivatives of reductions.
-//!
-//! A reduction fans one output back out to many inputs, so its derivative is a
-//! broadcast: every element that was summed gets the same gradient.
+//! Derivatives of reductions, which broadcast the gradient back out.
use crate::autograd::Backward;
use crate::tensor::shape;
use crate::tensor::Tensor;
-/// Every element contributed 1 to the sum, so each gets the whole gradient.
+/// Every element gets the whole gradient.
pub struct SumBackward {
pub shape: Vec,
}
@@ -21,7 +18,7 @@ impl Backward for SumBackward {
}
}
-/// Same as sum, divided by the number of elements averaged.
+/// Like sum, divided by the element count.
pub struct MeanBackward {
pub shape: Vec,
pub count: usize,
@@ -37,7 +34,7 @@ impl Backward for MeanBackward {
}
}
-/// Copy the gradient back along the axis that was collapsed.
+/// Copy the gradient back along the collapsed axis.
pub struct SumAxisBackward {
pub shape: Vec,
pub axis: usize,
diff --git a/src/autograd/ops/unary.rs b/src/autograd/ops/unary.rs
index 29a0178..4d1a283 100644
--- a/src/autograd/ops/unary.rs
+++ b/src/autograd/ops/unary.rs
@@ -1,12 +1,10 @@
-//! Derivatives of element-wise maths.
-//!
-//! Each rule multiplies the incoming gradient by a local derivative. Where that
-//! derivative is a simple function of the saved input, `elementwise` builds it.
+//! Derivatives of element-wise maths. Most multiply the gradient by a local
+//! derivative, which `elementwise` builds on the host.
use crate::autograd::Backward;
use crate::tensor::Tensor;
-/// `g · f(x)`, with `f` evaluated on the host and moved to the gradient's device.
+/// `g · f(x)`, with `f` evaluated on the host.
pub(crate) fn elementwise(
grad: &Tensor,
input: &Tensor,
@@ -28,7 +26,7 @@ impl Backward for NegBackward {
}
}
-/// `d(eˣ) = g·eˣ` — reuses the forward output.
+/// `d(eˣ) = g·eˣ`.
pub struct ExpBackward {
pub output: Tensor,
}
@@ -56,7 +54,7 @@ impl Backward for LogBackward {
}
}
-/// `d(√x) = g / (2√x)` — reuses the forward output.
+/// `d(√x) = g / (2√x)`.
pub struct SqrtBackward {
pub output: Tensor,
}
@@ -102,8 +100,7 @@ impl Backward for PowBackward {
}
}
-/// Clamping passes the gradient only where the input was inside the range;
-/// elsewhere the output was constant.
+/// Passes the gradient only where the input was inside the range.
pub struct ClampBackward {
pub input: Tensor,
pub lo: f32,
diff --git a/src/autograd/ops/view.rs b/src/autograd/ops/view.rs
index 8a57035..acc86c2 100644
--- a/src/autograd/ops/view.rs
+++ b/src/autograd/ops/view.rs
@@ -1,7 +1,4 @@
-//! Derivatives of the rearranging ops.
-//!
-//! Every one of these is the inverse rearrangement: reshape back, permute back,
-//! sum where elements were duplicated, pad with zeros where they were dropped.
+//! Derivatives of the rearranging ops, each the inverse rearrangement.
use crate::autograd::Backward;
use crate::tensor::shape;
@@ -42,7 +39,7 @@ impl Backward for PermuteBackward {
}
}
-/// Expanding copied one element to many, so its gradient is their sum.
+/// The gradient of a copy is the sum over the copies.
pub struct ExpandBackward {
pub shape: Vec,
}
@@ -56,7 +53,7 @@ impl Backward for ExpandBackward {
}
}
-/// Elements outside the slice did not reach the output, so they get zero.
+/// Elements outside the slice get zero.
pub struct NarrowBackward {
pub shape: Vec,
pub axis: usize,
diff --git a/src/cuda/buffer.rs b/src/cuda/buffer.rs
index a2ab9f0..1e35a0a 100644
--- a/src/cuda/buffer.rs
+++ b/src/cuda/buffer.rs
@@ -11,27 +11,21 @@ const F32: usize = std::mem::size_of::();
/// An owned block of `len` floats in GPU memory.
///
-/// Dropping it returns the block to the size-keyed free list rather than calling
-/// `cudaFree`, so the many short-lived temporaries a training step creates cost
-/// a hash lookup instead of a ~3 ms driver round trip.
+/// Dropping returns the block to a free list instead of calling `cudaFree`, which is slow.
pub struct CudaBuffer {
ptr: *mut f32,
len: usize,
}
-// GPU pointers are not tied to the thread that allocated them, and the CUDA
-// runtime serialises memory calls internally.
+// Device pointers are not tied to a thread, and the runtime serialises memory calls.
unsafe impl Send for CudaBuffer {}
unsafe impl Sync for CudaBuffer {}
impl CudaBuffer {
- /// Allocate `len` floats, reusing a cached block when one is free.
+ /// Allocates `len` floats, reusing a cached block when one is free.
///
- /// On failure the free list is drained and the allocation retried once. The
- /// cache holds blocks of every size that has ever been asked for, so a run
- /// that changes shape — a final short batch, a longer sequence — can be
- /// holding plenty of memory in the wrong size classes. Returning them to the
- /// driver turns many would-be out-of-memory failures into a pause.
+ /// On failure the cache is drained and the allocation retried once, since a changed shape
+ /// (a short last batch, a longer sequence) can leave memory parked in the wrong sizes.
pub fn new(len: usize) -> Result {
if let Some(ptr) = cache::take(len) {
return Ok(CudaBuffer { ptr, len });
@@ -45,7 +39,7 @@ impl CudaBuffer {
}
}
- /// Allocate `len` floats set to zero.
+ /// Allocates `len` zeroed floats.
pub fn zeros(len: usize) -> Result {
let buf = Self::new(len)?;
check(
@@ -55,14 +49,14 @@ impl CudaBuffer {
Ok(buf)
}
- /// Allocate and upload `data`.
+ /// Allocates a buffer holding `data`.
pub fn from_slice(data: &[f32]) -> Result {
let buf = Self::new(data.len())?;
buf.upload(data)?;
Ok(buf)
}
- /// Copy host memory into this buffer. `data` must not be longer than the buffer.
+ /// Copies host memory in. `data` must fit.
pub fn upload(&self, data: &[f32]) -> Result<()> {
assert!(
data.len() <= self.len,
@@ -76,7 +70,7 @@ impl CudaBuffer {
)
}
- /// Download the whole buffer to the host.
+ /// Copies the whole buffer to the host.
pub fn to_vec(&self) -> Result> {
let mut out = vec![0.0f32; self.len];
check(
@@ -86,7 +80,7 @@ impl CudaBuffer {
Ok(out)
}
- /// Copy `src` into this buffer on-device.
+ /// Copies `src` into this buffer on the device.
pub fn copy_from(&self, src: &CudaBuffer) -> Result<()> {
assert!(
src.len <= self.len,
@@ -164,15 +158,12 @@ fn raw_alloc(len: usize) -> Result<*mut f32> {
Ok(ptr)
}
-/// Free list of GPU blocks, keyed by element count.
-///
-/// Every kernel runs on the default stream, which is ordered: a kernel launched
-/// after another on the same stream sees its writes. So a block freed by one op
-/// is safe to hand to the next.
+/// Free list of blocks keyed by element count. Safe to reuse right away because every kernel
+/// runs on the ordered default stream.
mod cache {
use super::{ffi, HashMap, Mutex};
- /// Spare blocks kept per size before we hand memory back to the driver.
+ /// Spare blocks kept per size before memory goes back to the driver.
const MAX_PER_SIZE: usize = 32;
struct Block(*mut f32);
@@ -203,7 +194,6 @@ mod cache {
}
}
- /// Hand every cached block back to the driver.
pub(super) fn drain() {
let Ok(mut guard) = FREE_LIST.lock() else {
return;
@@ -218,10 +208,7 @@ mod cache {
}
}
-/// Release every cached GPU block back to the driver.
-///
-/// Only useful when another library needs the memory — allocation reuses cached
-/// blocks automatically.
+/// Frees every cached block. Only needed when another library wants the memory.
pub fn empty_cache() {
cache::drain();
}
diff --git a/src/cuda/ffi.rs b/src/cuda/ffi.rs
index 10180f9..75f5b94 100644
--- a/src/cuda/ffi.rs
+++ b/src/cuda/ffi.rs
@@ -1,23 +1,17 @@
-//! Every raw declaration of a symbol in `cuda/kernels.cu`, in one place.
+//! Raw `extern "C"` declarations for `cuda/kernels.cu`. Only [`super::kernels`] calls these.
//!
-//! Nothing outside this module calls these directly — [`super::kernels`] wraps
-//! them in safe, `Result`-returning functions. When you add a kernel to
-//! `kernels.cu`, add its signature here and a wrapper there.
-//!
-//! Every function returns `0` on success and non-zero on failure.
+//! Every function returns 0 on success.
#![allow(dead_code)]
use std::ffi::c_int;
extern "C" {
- // ── Context ──────────────────────────────────────────────────────────────
pub fn fastnn_cuda_init(device_id: c_int) -> c_int;
pub fn fastnn_cuda_device_count() -> c_int;
pub fn fastnn_cuda_synchronize() -> c_int;
pub fn fastnn_cuda_get_memory_info(free: *mut usize, total: *mut usize);
- // ── Memory ───────────────────────────────────────────────────────────────
pub fn fastnn_cuda_malloc(ptr: *mut *mut f32, bytes: usize) -> c_int;
pub fn fastnn_cuda_free(ptr: *mut f32) -> c_int;
pub fn fastnn_cuda_memcpy_h2d(dst: *mut f32, src: *const f32, bytes: usize) -> c_int;
@@ -26,7 +20,6 @@ extern "C" {
pub fn fastnn_cuda_memset(ptr: *mut f32, value: c_int, bytes: usize) -> c_int;
pub fn fastnn_cuda_fill(data: *mut f32, value: f32, n: usize) -> c_int;
- // ── Element-wise: binary, scalar, unary ──────────────────────────────────
pub fn fastnn_cuda_add(a: *const f32, b: *const f32, out: *mut f32, n: usize) -> c_int;
pub fn fastnn_cuda_sub(a: *const f32, b: *const f32, out: *mut f32, n: usize) -> c_int;
pub fn fastnn_cuda_mul(a: *const f32, b: *const f32, out: *mut f32, n: usize) -> c_int;
@@ -43,7 +36,6 @@ extern "C" {
pub fn fastnn_cuda_log(a: *const f32, out: *mut f32, n: usize) -> c_int;
pub fn fastnn_cuda_clamp(a: *const f32, lo: f32, hi: f32, out: *mut f32, n: usize) -> c_int;
- // ── Activations ──────────────────────────────────────────────────────────
pub fn fastnn_cuda_relu(x: *const f32, out: *mut f32, n: usize) -> c_int;
pub fn fastnn_cuda_relu_backward(
grad: *const f32,
@@ -81,7 +73,6 @@ extern "C" {
) -> c_int;
pub fn fastnn_cuda_leaky_relu(x: *const f32, slope: f32, out: *mut f32, n: usize) -> c_int;
- // ── Softmax (over the last dimension) ────────────────────────────────────
pub fn fastnn_cuda_softmax(x: *const f32, out: *mut f32, rows: c_int, cols: c_int) -> c_int;
pub fn fastnn_cuda_softmax_backward(
grad: *const f32,
@@ -93,7 +84,6 @@ extern "C" {
pub fn fastnn_cuda_log_softmax(x: *const f32, out: *mut f32, rows: c_int, cols: c_int)
-> c_int;
- // ── LayerNorm (fused; saves mean and 1/std for the backward pass) ────────
pub fn fastnn_cuda_layer_norm_forward(
x: *const f32,
gamma: *const f32,
@@ -118,7 +108,6 @@ extern "C" {
cols: c_int,
) -> c_int;
- // ── Embedding ────────────────────────────────────────────────────────────
pub fn fastnn_cuda_embedding_forward(
ids: *const c_int,
weight: *const f32,
@@ -135,7 +124,6 @@ extern "C" {
vocab: c_int,
) -> c_int;
- // ── Convolution lowering ─────────────────────────────────────────────────
pub fn fastnn_cuda_im2col(
input: *const f32,
cols: *mut f32,
@@ -173,7 +161,6 @@ extern "C" {
out_w: c_int,
) -> c_int;
- // ── GEMM (cuBLAS). `nt`/`tn` transpose in place, with no staging buffer. ─
pub fn fastnn_cuda_matmul(
a: *const f32,
b: *const f32,
@@ -233,7 +220,6 @@ extern "C" {
batch: c_int,
) -> c_int;
- // ── Layout ───────────────────────────────────────────────────────────────
pub fn fastnn_cuda_transpose(x: *const f32, out: *mut f32, rows: c_int, cols: c_int) -> c_int;
pub fn fastnn_cuda_transpose_batched(
x: *const f32,
@@ -255,7 +241,6 @@ extern "C" {
numel: c_int,
) -> c_int;
- // ── Reductions ───────────────────────────────────────────────────────────
pub fn fastnn_cuda_sum(x: *const f32, out: *mut f32, n: usize) -> c_int;
pub fn fastnn_cuda_mean(x: *const f32, out: *mut f32, n: usize) -> c_int;
pub fn fastnn_cuda_max(x: *const f32, out: *mut f32, n: usize) -> c_int;
diff --git a/src/cuda/kernels.rs b/src/cuda/kernels.rs
index bb747d3..117a42a 100644
--- a/src/cuda/kernels.rs
+++ b/src/cuda/kernels.rs
@@ -1,7 +1,4 @@
-//! Safe wrappers around the CUDA kernels.
-//!
-//! Each function allocates its output buffer, launches the kernel, and turns a
-//! non-zero status into an [`Error::Cuda`]. Callers never touch raw pointers.
+//! Safe wrappers over [`super::ffi`]: allocate the output, launch, map a bad status to [`Error::Cuda`].
use std::ffi::c_int;
@@ -9,7 +6,7 @@ use super::buffer::CudaBuffer;
use super::ffi::{self, BinaryKernel, ScalarKernel, UnaryKernel};
use crate::error::{Error, Result};
-/// Run `kernel`, returning its freshly allocated output of `out_len` floats.
+/// Runs `kernel` on a fresh output of `out_len` floats.
fn launch(
out_len: usize,
name: &'static str,
@@ -23,8 +20,6 @@ fn launch(
Ok(out)
}
-// ── Element-wise ─────────────────────────────────────────────────────────────
-
pub fn binary(
a: &CudaBuffer,
b: &CudaBuffer,
@@ -60,8 +55,6 @@ pub fn leaky_relu(a: &CudaBuffer, slope: f32, n: usize) -> Result {
})
}
-// ── Softmax ──────────────────────────────────────────────────────────────────
-
pub fn softmax(x: &CudaBuffer, rows: usize, cols: usize) -> Result {
launch(rows * cols, "softmax", |out| unsafe {
ffi::fastnn_cuda_softmax(x.as_ptr(), out.as_mut_ptr(), rows as c_int, cols as c_int)
@@ -91,8 +84,6 @@ pub fn softmax_backward(
})
}
-// ── GEMM ─────────────────────────────────────────────────────────────────────
-
/// `C[m,n] = A[m,k] · B[k,n]`
pub fn matmul(a: &CudaBuffer, b: &CudaBuffer, m: usize, n: usize, k: usize) -> Result {
launch(m * n, "matmul", |out| unsafe {
@@ -112,7 +103,7 @@ pub fn matmul(a: &CudaBuffer, b: &CudaBuffer, m: usize, n: usize, k: usize) -> R
})
}
-/// `C[m,n] = A[m,k] · B[n,k]ᵀ` — B is read transposed, no staging buffer.
+/// `C[m,n] = A[m,k] · B[n,k]ᵀ`
pub fn matmul_nt(
a: &CudaBuffer,
b: &CudaBuffer,
@@ -132,7 +123,7 @@ pub fn matmul_nt(
})
}
-/// `C[m,n] = A[k,m]ᵀ · B[k,n]` — A is read transposed, no staging buffer.
+/// `C[m,n] = A[k,m]ᵀ · B[k,n]`
pub fn matmul_tn(
a: &CudaBuffer,
b: &CudaBuffer,
@@ -217,8 +208,6 @@ pub fn matmul_batched_tn(
})
}
-// ── Layout ───────────────────────────────────────────────────────────────────
-
pub fn transpose_2d(x: &CudaBuffer, rows: usize, cols: usize) -> Result {
launch(rows * cols, "transpose", |out| unsafe {
ffi::fastnn_cuda_transpose(x.as_ptr(), out.as_mut_ptr(), rows as c_int, cols as c_int)
@@ -242,9 +231,8 @@ pub fn transpose_batched(
})
}
-/// Gather `numel` elements: output index `i` reads the input element whose
-/// coordinates come from decoding `i` against `out_strides` and re-encoding with
-/// `in_strides` permuted by `perm`. Stride 0 broadcasts a dimension.
+/// Gathers `numel` elements by decoding each output index with `out_strides` and re-encoding with
+/// `in_strides` permuted by `perm`. A stride of 0 broadcasts.
pub fn permute_nd(
x: &CudaBuffer,
out_strides: &[usize],
@@ -268,8 +256,6 @@ pub fn permute_nd(
})
}
-// ── Reductions ───────────────────────────────────────────────────────────────
-
pub fn sum(x: &CudaBuffer, n: usize) -> Result {
launch(1, "sum", |out| unsafe {
ffi::fastnn_cuda_sum(x.as_ptr(), out.as_mut_ptr(), n)
@@ -309,8 +295,6 @@ pub fn sum_axis(x: &CudaBuffer, shape: &[usize], axis: usize) -> Result Result {
let as_f32 = unsafe { std::slice::from_raw_parts(ids.as_ptr() as *const f32, ids.len()) };
CudaBuffer::from_slice(as_f32)
@@ -412,7 +393,7 @@ pub fn embedding_backward(
dim: usize,
vocab: usize,
) -> Result {
- // The kernel scatter-adds, so the accumulator must start at zero.
+ // the kernel scatter-adds
let out = CudaBuffer::zeros(vocab * dim)?;
let code = unsafe {
ffi::fastnn_cuda_embedding_backward(
@@ -432,9 +413,7 @@ pub fn embedding_backward(
Ok(out)
}
-// ── Convolution lowering ─────────────────────────────────────────────────────
-
-/// Unfold `[n, c, h, w]` into `[n, c·kh·kw, out_h·out_w]` columns on the device.
+/// Unfolds `[n, c, h, w]` into `[n, c·kh·kw, out_h·out_w]` columns.
pub fn im2col(
input: &CudaBuffer,
(n, c, h, w): (usize, usize, usize, usize),
@@ -466,7 +445,7 @@ pub fn im2col(
})
}
-/// Fold columns back into `[n, c, h, w]`, summing where windows overlapped.
+/// Folds columns back into `[n, c, h, w]`, summing overlaps.
pub fn col2im(
cols: &CudaBuffer,
(n, c, h, w): (usize, usize, usize, usize),
diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs
index 3b38df5..fa58403 100644
--- a/src/cuda/mod.rs
+++ b/src/cuda/mod.rs
@@ -1,11 +1,6 @@
-//! GPU layer: device setup, memory, and the kernel bindings.
+//! GPU layer: device setup, memory ([`buffer`]), raw bindings ([`ffi`]) and safe wrappers ([`kernels`]).
//!
-//! - [`ffi`] — raw `extern "C"` declarations, one per symbol in `cuda/kernels.cu`
-//! - [`kernels`] — safe wrappers that allocate outputs and check status codes
-//! - [`buffer`] — [`CudaBuffer`], RAII GPU memory over a caching allocator
-//!
-//! In a `--no-default-features` build, `cuda/stubs.c` supplies the symbols so
-//! everything links; [`device_count`] then reports 0 and CPU paths run instead.
+//! Without the `cuda` feature, `cuda/stubs.c` stands in for the kernels and [`device_count`] is 0.
pub mod buffer;
pub mod ffi;
@@ -17,7 +12,7 @@ use std::sync::OnceLock;
use crate::error::{Error, Result};
-/// Initialise the CUDA context on `device`. Idempotent; safe to call per tensor.
+/// Sets up the CUDA context on `device`. Only the first call does anything.
pub fn init(device: usize) -> Result<()> {
static INIT: OnceLock> = OnceLock::new();
INIT.get_or_init(|| {
@@ -32,12 +27,12 @@ pub fn init(device: usize) -> Result<()> {
.map_err(Error::Cuda)
}
-/// Number of CUDA devices visible to the process. 0 in a CPU-only build.
+/// Number of visible CUDA devices, 0 in a CPU-only build.
pub fn device_count() -> usize {
unsafe { ffi::fastnn_cuda_device_count().max(0) as usize }
}
-/// Block until every queued kernel has finished.
+/// Waits for every queued kernel to finish.
pub fn synchronize() -> Result<()> {
match unsafe { ffi::fastnn_cuda_synchronize() } {
0 => Ok(()),
diff --git a/src/data/dataset.rs b/src/data/dataset.rs
index 8982d0b..0675639 100644
--- a/src/data/dataset.rs
+++ b/src/data/dataset.rs
@@ -1,13 +1,11 @@
-//! The dataset trait and simple in-memory implementations.
+//! The dataset trait and an in-memory implementation.
use crate::tensor::{shape, Tensor};
/// A finite, indexable collection of `(input, target)` pairs.
///
-/// Items are written into caller-provided slices rather than returned as
-/// tensors. A [`DataLoader`](super::DataLoader) allocates one buffer per batch
-/// and fills it, so batching a dataset of 60,000 images does not build 60,000
-/// throwaway tensors.
+/// Items are written into slices so a [`DataLoader`](super::DataLoader) can fill
+/// one buffer per batch instead of building a tensor per item.
pub trait Dataset: Send + Sync {
/// Number of items.
fn len(&self) -> usize;
diff --git a/src/data/loader.rs b/src/data/loader.rs
index dc19063..45f5b19 100644
--- a/src/data/loader.rs
+++ b/src/data/loader.rs
@@ -12,7 +12,7 @@ pub struct Batch {
}
impl Batch {
- /// How many items are in this batch.
+ /// Items in this batch.
pub fn len(&self) -> usize {
self.inputs.dim(0)
}
@@ -21,9 +21,9 @@ impl Batch {
self.len() == 0
}
- /// Targets read as class indices — the form [`cross_entropy`](crate::nn::cross_entropy) takes.
+ /// Targets as class indices, for [`cross_entropy`](crate::nn::cross_entropy).
///
- /// Panics if a target is not a finite, non-negative integer representable as `usize`.
+ /// Panics if a target is not a finite, non-negative integer that fits in `usize`.
pub fn labels(&self) -> Vec {
self.targets
.to_vec()
@@ -48,8 +48,7 @@ impl Batch {
/// }
/// ```
///
-/// A fresh shuffle happens each time [`iter`](DataLoader::iter) is called, so
-/// every epoch sees a different order.
+/// Each call to [`iter`](DataLoader::iter) reshuffles.
pub struct DataLoader<'a> {
dataset: &'a dyn Dataset,
batch_size: usize,
@@ -70,13 +69,13 @@ impl<'a> DataLoader<'a> {
}
}
- /// Reorder items each epoch. On for training, off for evaluation.
+ /// Reorder items each epoch.
pub fn shuffle(mut self, shuffle: bool) -> DataLoader<'a> {
self.shuffle = shuffle;
self
}
- /// Skip the final short batch. Useful when a layer needs a fixed batch size.
+ /// Skip the final short batch.
pub fn drop_last(mut self, drop_last: bool) -> DataLoader<'a> {
self.drop_last = drop_last;
self
@@ -88,7 +87,7 @@ impl<'a> DataLoader<'a> {
self
}
- /// Number of batches one pass will yield.
+ /// Batches in one pass.
pub fn batches(&self) -> usize {
let n = self.dataset.len();
if self.drop_last {
@@ -124,7 +123,7 @@ impl<'a> IntoIterator for &DataLoader<'a> {
}
}
-/// One pass over a shuffled index order.
+/// One pass over the data.
pub struct Batches<'a> {
dataset: &'a dyn Dataset,
order: Vec,
@@ -175,8 +174,7 @@ impl Iterator for Batches<'_> {
}
}
-/// Prepend the batch dimension. A `[1]` item shape collapses to `[batch]`, which
-/// is what a vector of class labels should look like.
+/// Prepend the batch dimension. A `[1]` item shape becomes `[batch]`, like a label vector.
fn batched(data: Vec, size: usize, item_shape: &[usize]) -> Tensor {
if item_shape == [1] {
return Tensor::from_vec(data, &[size]);
diff --git a/src/data/mnist.rs b/src/data/mnist.rs
index b2129e0..aef3fdb 100644
--- a/src/data/mnist.rs
+++ b/src/data/mnist.rs
@@ -1,7 +1,5 @@
-//! MNIST handwritten digits.
-//!
-//! Downloads the four IDX files on first use and caches them under
-//! `$FASTNN_DATA_DIR/mnist` (or `~/.fastnn/datasets/mnist`).
+//! MNIST digits, downloaded on first use into `$FASTNN_DATA_DIR/mnist` or
+//! `~/.fastnn/datasets/mnist`.
use std::fs;
use std::io::{Read, Write};
@@ -15,6 +13,10 @@ use super::dataset::Dataset;
const MIRROR: &str = "https://ossci-datasets.s3.amazonaws.com/mnist";
const IMAGE_MAGIC: u32 = 0x0000_0803;
const LABEL_MAGIC: u32 = 0x0000_0801;
+const SIDE: usize = 28;
+const PIXELS: usize = SIDE * SIDE;
+const IMAGE_HEADER: usize = 16;
+const LABEL_HEADER: usize = 8;
/// Which half of the dataset to load.
#[derive(Clone, Copy, Debug)]
@@ -32,9 +34,7 @@ impl Split {
}
}
-/// 28×28 grayscale digits with labels 0–9.
-///
-/// Images are `[1, 28, 28]` scaled to `[0, 1]`; targets are class indices.
+/// 28x28 grayscale digits scaled to `[0, 1]`, with class index targets.
///
/// ```no_run
/// use fastnn::prelude::*;
@@ -76,12 +76,12 @@ impl Mnist {
})
}
- /// Every image as one `[N, 1, 28, 28]` tensor.
+ /// All images as one `[N, 1, 28, 28]` tensor.
pub fn images(&self) -> Tensor {
- Tensor::from_vec(self.images.clone(), &[self.len, 1, 28, 28])
+ Tensor::from_vec(self.images.clone(), &[self.len, 1, SIDE, SIDE])
}
- /// Every label as a class index.
+ /// All labels as class indices.
pub fn labels(&self) -> Vec {
self.labels.iter().map(|&v| v as usize).collect()
}
@@ -93,7 +93,7 @@ impl Dataset for Mnist {
}
fn input_shape(&self) -> Vec {
- vec![1, 28, 28]
+ vec![1, SIDE, SIDE]
}
fn target_shape(&self) -> Vec {
@@ -101,8 +101,8 @@ impl Dataset for Mnist {
}
fn write(&self, index: usize, input: &mut [f32], target: &mut [f32]) {
- let start = index * 28 * 28;
- input.copy_from_slice(&self.images[start..start + 28 * 28]);
+ let start = index * PIXELS;
+ input.copy_from_slice(&self.images[start..start + PIXELS]);
target[0] = self.labels[index];
}
}
@@ -119,7 +119,7 @@ pub fn default_cache_dir() -> PathBuf {
.join("mnist")
}
-/// Fetch and decompress any of `files` that are not already cached.
+/// Download and decompress any of `files` not already cached.
fn download_missing(cache: &Path, files: &[&str]) -> Result<()> {
fs::create_dir_all(cache)?;
for name in files {
@@ -152,16 +152,18 @@ fn parse_images(bytes: &[u8]) -> Result<(Vec, usize)> {
let header = read_header(bytes, 4, IMAGE_MAGIC, "images")?;
let (count, rows, cols) = (header[1], header[2], header[3]);
- let expected = 16 + count * rows * cols;
+ let expected = IMAGE_HEADER + count * rows * cols;
if bytes.len() != expected {
return Err(Error::Dataset(format!(
"mnist images: {} bytes, expected {expected}",
bytes.len()
)));
}
- // Scale to [0, 1]; unnormalized 0–255 inputs would saturate the first layer.
Ok((
- bytes[16..].iter().map(|&b| b as f32 / 255.0).collect(),
+ bytes[IMAGE_HEADER..]
+ .iter()
+ .map(|&b| b as f32 / 255.0)
+ .collect(),
count,
))
}
@@ -171,14 +173,17 @@ fn parse_labels(bytes: &[u8]) -> Result<(Vec, usize)> {
let header = read_header(bytes, 2, LABEL_MAGIC, "labels")?;
let count = header[1];
- if bytes.len() != 8 + count {
+ if bytes.len() != LABEL_HEADER + count {
return Err(Error::Dataset(format!(
"mnist labels: {} bytes, expected {}",
bytes.len(),
- 8 + count
+ LABEL_HEADER + count
)));
}
- Ok((bytes[8..].iter().map(|&b| b as f32).collect(), count))
+ Ok((
+ bytes[LABEL_HEADER..].iter().map(|&b| b as f32).collect(),
+ count,
+ ))
}
/// Read `words` big-endian u32s and check the magic number.
diff --git a/src/data/mod.rs b/src/data/mod.rs
index c65fd38..2b6f322 100644
--- a/src/data/mod.rs
+++ b/src/data/mod.rs
@@ -1,7 +1,4 @@
-//! Datasets and batching.
-//!
-//! Implement [`Dataset`] for your own data, then wrap it in a [`DataLoader`] to
-//! get shuffling, batching, and device placement.
+//! Datasets, batching, and MNIST.
pub mod dataset;
pub mod loader;
diff --git a/src/error.rs b/src/error.rs
index 9d800cd..9f6c999 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,13 +1,10 @@
-//! Error type for the fallible boundaries of the library.
+//! The error type for I/O, downloads, checkpoints, and CUDA.
//!
-//! Tensor math does *not* return `Result`. A shape or device mismatch is a bug in
-//! the calling code, like indexing past the end of a slice, so it panics with a
-//! message naming both shapes. `Result` is reserved for things that can fail at
-//! runtime through no fault of the code: I/O, downloads, and CUDA.
+//! Shape and device mistakes in tensor maths panic instead.
use std::fmt;
-/// Everything that can go wrong at a fallible boundary.
+/// Everything that can fail at runtime.
#[derive(Debug)]
pub enum Error {
/// A CUDA driver, runtime, or kernel call failed.
@@ -22,12 +19,11 @@ pub enum Error {
Dataset(String),
/// Shapes or devices did not line up.
///
- /// Only the [`try_*`](crate::tensor::Tensor::try_matmul) ops produce this.
- /// The ordinary ops panic instead — see the module docs.
+ /// Only the [`try_*`](crate::tensor::Tensor::try_matmul) ops return this.
Shape(String),
}
-/// Shorthand for results carrying a fastnn [`Error`].
+/// Result alias for fastnn's [`Error`].
pub type Result = std::result::Result;
impl fmt::Display for Error {
diff --git a/src/lib.rs b/src/lib.rs
index 7bc9f5f..927979b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,6 +1,6 @@
//! # FastNN
//!
-//! A deep learning library in Rust, with CUDA kernels for the parts that matter.
+//! A deep learning library in Rust with optional CUDA kernels.
//!
//! ```no_run
//! use fastnn::prelude::*;
@@ -31,42 +31,14 @@
//! # }
//! ```
//!
-//! ## How it fits together
+//! Modules: [`tensor`], [`autograd`], [`nn`], [`optim`], [`data`], [`serialize`], [`cuda`].
//!
-//! ```text
-//! data Dataset ─► DataLoader ─► Batch
-//! │
-//! nn Module ──► forward ────────┤ layers hold Param handles
-//! ▼
-//! tensor Tensor ops on CPU or CUDA ─┴─► loss
-//! │
-//! autograd loss.backward() walks the graph the ops built,
-//! depositing gradients in each Param's slot
-//! │
-//! optim opt.step() reads those slots and updates in place
-//! ```
-//!
-//! Each layer is one module directory and is documented where it lives:
-//!
-//! - [`tensor`] — the array type, its ops, and CPU/CUDA dispatch
-//! - [`autograd`] — the graph the ops build and the reverse pass over it
-//! - [`nn`] — layers, [`Param`](nn::Param) handles, and losses
-//! - [`optim`] — optimizers and learning-rate schedules
-//! - [`data`] — datasets, batching, MNIST
-//! - [`serialize`] — checkpoints
-//! - [`cuda`] — GPU memory and the kernel bindings
-//!
-//! ## Errors
-//!
-//! Tensor maths panics on shape and device mistakes — those are bugs in the
-//! calling code, like indexing past the end of a slice. [`Result`] is for things
-//! that genuinely fail at runtime: I/O, dataset downloads, and CUDA.
-//!
-//! ## Devices
+//! Tensor maths panics on shape and device mistakes. [`Result`] is for I/O,
+//! dataset downloads, and CUDA.
//!
-//! Everything starts on the CPU. [`Module::to_device`](nn::Module::to_device)
-//! moves a model, [`DataLoader::to_device`](data::DataLoader::to_device) moves
-//! its batches, and mixing the two panics rather than copying silently.
+//! Everything starts on the CPU. Move the model with
+//! [`Module::to_device`](nn::Module::to_device) and the batches with
+//! [`DataLoader::to_device`](data::DataLoader::to_device).
//!
//! ```no_run
//! # use fastnn::prelude::*;
@@ -89,7 +61,7 @@ pub mod tensor;
pub use error::{Error, Result};
-/// Everything you need for a training script.
+/// The usual imports for a training script.
///
/// ```
/// use fastnn::prelude::*;
diff --git a/src/nn/activation.rs b/src/nn/activation.rs
index e4b9f3b..e0af864 100644
--- a/src/nn/activation.rs
+++ b/src/nn/activation.rs
@@ -1,13 +1,9 @@
-//! Activations as layers, for use inside [`Sequential`](super::Sequential).
-//!
-//! Each one is a thin wrapper over the corresponding [`Tensor`] method — call
-//! that directly when you are writing a `forward` by hand.
+//! Activations as layers, thin wrappers over the [`Tensor`] methods.
use crate::tensor::Tensor;
use super::module::Module;
-/// Declare a stateless activation layer that forwards to one tensor method.
macro_rules! activation {
($(#[$doc:meta])* $name:ident => $method:ident) => {
$(#[$doc])*
@@ -27,17 +23,18 @@ activation!(/// `1 / (1 + e^-x)`.
Sigmoid => sigmoid);
activation!(/// Hyperbolic tangent.
Tanh => tanh);
-activation!(/// Exact GELU — the transformer default.
+activation!(/// Exact GELU, the transformer default.
GELU => gelu);
activation!(/// SiLU / swish, `x · sigmoid(x)`.
SiLU => silu);
activation!(/// Softmax over the last dimension.
///
- /// Not needed before [`cross_entropy`](super::cross_entropy), which
- /// applies its own and is more stable for it.
+ /// Skip it before [`cross_entropy`](super::cross_entropy), which has its own.
Softmax => softmax);
-/// ReLU with a non-zero slope for negative inputs, so units cannot go fully dead.
+const DEFAULT_SLOPE: f32 = 0.01;
+
+/// ReLU with a small slope for negative inputs.
pub struct LeakyReLU {
pub slope: f32,
}
@@ -50,7 +47,9 @@ impl LeakyReLU {
impl Default for LeakyReLU {
fn default() -> Self {
- LeakyReLU { slope: 0.01 }
+ LeakyReLU {
+ slope: DEFAULT_SLOPE,
+ }
}
}
diff --git a/src/nn/attention.rs b/src/nn/attention.rs
index 0982afb..8943dd0 100644
--- a/src/nn/attention.rs
+++ b/src/nn/attention.rs
@@ -11,10 +11,11 @@ use super::linear::Linear;
use super::module::{scoped, Module};
use super::param::Param;
-/// Attention over `[batch, sequence, embed_dim]`.
-///
-/// Heads are not separate tensors: one projection produces all of them, and the
-/// head axis is folded into the batch so a single batched GEMM covers every head.
+/// Large negative added to masked scores, so softmax gives them zero weight.
+const MASKED: f32 = -1e9;
+
+/// Attention over `[batch, sequence, embed_dim]`. The head axis is folded into the
+/// batch, so one batched GEMM covers every head.
pub struct MultiHeadAttention {
pub query: Linear,
pub key: Linear,
@@ -42,25 +43,18 @@ impl MultiHeadAttention {
dropout: Dropout::new(dropout),
heads,
head_dim,
- // Without this, dot products grow with head_dim and softmax saturates.
+ // dot products grow with head_dim and would saturate the softmax
scale: 1.0 / (head_dim as f32).sqrt(),
}
}
- /// Attend from `query` over `key`/`value`.
- ///
- /// With `causal` set, position `i` can only see positions `≤ i` — what a
- /// language model needs so it cannot read ahead.
+ /// Attend from `query` over `key`/`value`. With `causal`, position `i` sees only `≤ i`.
pub fn attend(&self, query: &Tensor, key: &Tensor, value: &Tensor, causal: bool) -> Tensor {
self.attend_masked(query, key, value, causal, None)
}
- /// [`attend`](Self::attend) with an optional key padding mask.
- ///
- /// `key_mask` is `[batch, kv_len]`: 1 where a key position is real, 0 where
- /// it is padding. Padded positions get the same additive `-1e9` treatment as
- /// the causal mask, so softmax gives them zero weight and batches of
- /// unequal-length sequences attend only to their own content.
+ /// [`attend`](Self::attend) with an optional padding mask, `[batch, kv_len]`,
+ /// 1 for real keys and 0 for padding.
pub fn attend_masked(
&self,
query: &Tensor,
@@ -77,11 +71,9 @@ impl MultiHeadAttention {
let k = self.split_heads(&self.key.forward(key), batch, kv_len);
let v = self.split_heads(&self.value.forward(value), batch, kv_len);
- // matmul_nt gives q · kᵀ without ever materialising kᵀ.
let mut scores = q.matmul_nt(&k).mul_scalar(self.scale);
if causal {
- // An additive -inf bias rather than an in-place overwrite: the mask
- // stays part of the graph, so gradients still flow through softmax.
+ // additive rather than an overwrite, so it stays in the graph
scores = scores.add(&causal_mask(lanes, q_len, kv_len, scores.device()));
}
if let Some(mask) = key_mask {
@@ -90,11 +82,10 @@ impl MultiHeadAttention {
&[batch, kv_len],
"key mask must be [batch, kv_len] = [{batch}, {kv_len}]"
);
- // 1/0 keep/pad → 0/-1e9, repeated per head so it broadcasts over
- // the [batch·heads, q_len, kv_len] scores.
+ // 1/0 becomes 0/MASKED, repeated per head to match the folded batch
let bias = mask
.add_scalar(-1.0)
- .mul_scalar(1e9)
+ .mul_scalar(-MASKED)
.reshape(&[batch as i64, 1, kv_len as i64])
.expand(&[batch, self.heads, kv_len])
.reshape(&[lanes as i64, 1, kv_len as i64]);
@@ -108,13 +99,8 @@ impl MultiHeadAttention {
.forward(&self.merge_heads(&attended, batch, q_len, embed_dim))
}
- /// Causal self-attention over `input`'s new tokens plus everything in
- /// `cache`, extending the cache as it goes.
- ///
- /// Inference-only: dropout is skipped and the cache stores detached
- /// tensors, so run it under [`no_grad`](crate::autograd::no_grad). An empty
- /// cache makes this a plain causal forward (the prefill); after that, feed
- /// one token at a time and only the new projections are computed.
+ /// Causal self-attention over the new tokens plus `cache`, which it extends.
+ /// Inference only: no dropout, run under [`no_grad`](crate::autograd::no_grad).
pub fn attend_cached(&self, input: &Tensor, cache: &mut KvCache) -> Tensor {
let (batch, new_len, embed_dim) = (input.dim(0), input.dim(1), input.dim(2));
@@ -125,8 +111,7 @@ impl MultiHeadAttention {
let kv_len = k.dim(1);
let mut scores = q.matmul_nt(k).mul_scalar(self.scale);
- // A single new token sits at the end of the sequence and may see all of
- // it, so only multi-token steps need masking.
+ // one new token may see everything, so only multi-token steps need a mask
if new_len > 1 {
scores = scores.add(&causal_mask(
batch * self.heads,
@@ -171,8 +156,7 @@ impl MultiHeadAttention {
}
impl Module for MultiHeadAttention {
- /// Self-attention with no mask. Use [`attend`](MultiHeadAttention::attend)
- /// for cross-attention or causal masking.
+ /// Unmasked self-attention. Use [`attend`](MultiHeadAttention::attend) for the rest.
fn forward(&self, input: &Tensor) -> Tensor {
self.attend(input, input, input, false)
}
@@ -195,18 +179,13 @@ impl Module for MultiHeadAttention {
}
thread_local! {
- /// Causal masks depend only on shape and device, and every layer of every
- /// step wants the same one — so build each exactly once per thread.
+ /// Every layer wants the same mask for a given shape, so build each once.
static MASKS: RefCell> =
RefCell::new(HashMap::new());
}
-/// `[lanes, q_len, kv_len]`, zero where attention is allowed and a large
-/// negative where it is not, so softmax drives those weights to zero.
-///
-/// When `kv_len > q_len` the queries are the *last* `q_len` positions of the
-/// sequence — a cached decoding step — so each sees the whole prefix plus
-/// itself. With equal lengths this is the ordinary causal triangle.
+/// `[lanes, q_len, kv_len]`, zero where attention is allowed and [`MASKED`] elsewhere.
+/// When `kv_len > q_len` the queries are the last `q_len` positions (a cached step).
fn causal_mask(lanes: usize, q_len: usize, kv_len: usize, device: Device) -> Tensor {
MASKS.with(|cache| {
cache
@@ -219,7 +198,7 @@ fn causal_mask(lanes: usize, q_len: usize, kv_len: usize, device: Device) -> Ten
let query_pos = (i / kv_len) % q_len;
let key_pos = i % kv_len;
if key_pos > query_pos + offset {
- -1e9
+ MASKED
} else {
0.0
}
diff --git a/src/nn/beam.rs b/src/nn/beam.rs
index 85db8dd..715f3ef 100644
--- a/src/nn/beam.rs
+++ b/src/nn/beam.rs
@@ -1,15 +1,6 @@
-//! Beam search decoding.
-//!
-//! Sampling explores; beam search commits. It keeps the `width` most probable
-//! partial sequences and extends each one step at a time, which finds higher
-//! total-probability outputs than greedy decoding — the best second token may
-//! hide behind the second-best first token. Width 1 *is* greedy decoding.
+//! Beam search decoding. Width 1 is greedy.
-/// Beam search over any autoregressive model.
-///
-/// The model appears only as a closure from a sequence to next-token logits,
-/// so this works with a [`TransformerStack`](super::TransformerStack), an RNN,
-/// or anything else that predicts one step at a time.
+/// Beam search over any model that maps a sequence to next-token logits.
///
/// ```
/// use fastnn::nn::BeamSearch;
@@ -28,7 +19,7 @@ pub struct BeamSearch {
eos: Option,
}
-/// One candidate sequence: the generated ids and their total log-probability.
+/// A candidate sequence and its total log-probability.
struct Beam {
generated: Vec,
log_prob: f32,
@@ -36,8 +27,7 @@ struct Beam {
}
impl Beam {
- /// Longer sequences accumulate more negative log-probability by sheer
- /// length; dividing by `len^penalty` keeps them comparable to short ones.
+ /// Length-normalised, since long sequences always have lower raw log-probability.
fn score(&self, length_penalty: f32) -> f32 {
self.log_prob / (self.generated.len().max(1) as f32).powf(length_penalty)
}
@@ -53,8 +43,7 @@ impl BeamSearch {
}
}
- /// Exponent on length when comparing beams. 1.0 is per-token average;
- /// below favours short outputs, above favours long ones.
+ /// Exponent on length when comparing beams. 1.0 is the per-token average.
pub fn length_penalty(mut self, penalty: f32) -> BeamSearch {
assert!(
penalty > 0.0,
@@ -70,11 +59,8 @@ impl BeamSearch {
self
}
- /// Generate up to `max_tokens` ids after `start`.
- ///
- /// `step` maps a full sequence (`start` plus everything generated so far)
- /// to the logits for the next token. Returns the generated ids of the best
- /// beam, without `start` and without the end-of-sequence token.
+ /// Generate up to `max_tokens` ids after `start`, without `start` or the eos token.
+ /// `step` maps the full sequence so far to next-token logits.
pub fn decode(
&self,
start: &[usize],
@@ -95,7 +81,7 @@ impl BeamSearch {
let mut candidates = Vec::new();
for beam in &beams {
if beam.finished {
- // A finished beam still competes on score, unchanged.
+ // finished beams keep competing on score
candidates.push(Beam {
generated: beam.generated.clone(),
log_prob: beam.log_prob,
@@ -108,8 +94,7 @@ impl BeamSearch {
sequence.extend_from_slice(&beam.generated);
let log_probs = log_softmax(&step(&sequence));
- // Only the `width` best continuations of this beam can survive
- // the global cut, so nothing else needs to be considered.
+ // only this beam's top `width` continuations can survive the global cut
let mut order: Vec = (0..log_probs.len()).collect();
order.sort_by(|&a, &b| log_probs[b].total_cmp(&log_probs[a]));
for &token in order.iter().take(self.width) {
diff --git a/src/nn/cache.rs b/src/nn/cache.rs
index 52b0259..d469d97 100644
--- a/src/nn/cache.rs
+++ b/src/nn/cache.rs
@@ -1,20 +1,11 @@
//! Cached keys and values for incremental decoding.
//!
-//! Generating token by token through [`Module::forward`](super::Module::forward)
-//! recomputes attention over the whole prefix for every new token. A cache keeps
-//! each layer's projected keys and values, so a step only computes projections
-//! for the tokens it has not seen — O(n) attention per token instead of O(n²).
-//!
-//! One [`KvCache`] serves one attention layer; a [`StackCache`] holds one per
-//! block of a [`TransformerStack`](super::TransformerStack). The cached tensors
-//! are detached, so holding a cache never pins an autograd graph — this is an
-//! inference path, and callers should run it under
+//! Detached on purpose, so a cache never pins a graph. Run it under
//! [`no_grad`](crate::autograd::no_grad).
use crate::tensor::Tensor;
-/// The keys and values one attention layer has already projected, each
-/// `[batch·heads, cached_len, head_dim]`.
+/// One attention layer's projected keys and values, each `[batch·heads, len, head_dim]`.
#[derive(Default)]
pub struct KvCache {
keys: Option,
@@ -41,9 +32,7 @@ impl KvCache {
self.values = None;
}
- /// Extend the cache with newly projected keys and values, returning the
- /// full cached tensors to attend over. Stores detached copies, so the graph
- /// that produced the projections is free to go.
+ /// Append new keys and values, returning the full cached tensors.
pub fn append(&mut self, keys: &Tensor, values: &Tensor) -> (&Tensor, &Tensor) {
self.keys = Some(match &self.keys {
Some(prior) => Tensor::cat(&[prior, &keys.detach()], 1),
@@ -57,7 +46,7 @@ impl KvCache {
}
}
-/// One [`KvCache`] per block of a transformer stack.
+/// One [`KvCache`] per block of a [`TransformerStack`](super::TransformerStack).
pub struct StackCache {
pub layers: Vec,
}
@@ -69,7 +58,7 @@ impl StackCache {
}
}
- /// How many positions are cached — the next token's position index.
+ /// How many positions are cached, which is the next token's position index.
pub fn len(&self) -> usize {
self.layers.first().map_or(0, KvCache::len)
}
diff --git a/src/nn/conv.rs b/src/nn/conv.rs
index 383f478..1241c1c 100644
--- a/src/nn/conv.rs
+++ b/src/nn/conv.rs
@@ -2,15 +2,11 @@
use crate::tensor::{Tensor, Window};
+use super::linear::uniform_bias;
use super::module::Module;
use super::param::Param;
-/// Convolution over `[N, C, H, W]`, implemented as unfold-then-multiply.
-///
-/// [`im2col`](Tensor::im2col) lays every sliding window out as a column, which
-/// turns the convolution into one matrix multiply against a `[out_channels,
-/// C·kh·kw]` weight. That is why this layer needs no gradient code of its own:
-/// the derivatives of `im2col` and `matmul` already compose into the right thing.
+/// Convolution over `[N, C, H, W]` as [`im2col`](Tensor::im2col) then one matmul.
pub struct Conv2d {
pub weight: Param,
pub bias: Option,
@@ -57,10 +53,6 @@ impl Conv2d {
}
/// Split the channels into `groups` independent convolutions.
- ///
- /// Each group sees only `in/groups` input channels and produces
- /// `out/groups` outputs, cutting parameters and compute by the group
- /// count. The weight is `[out, in/groups, kh, kw]`.
pub fn grouped(
in_channels: usize,
out_channels: usize,
@@ -78,9 +70,7 @@ impl Conv2d {
)
}
- /// One filter per channel — `groups == channels`, the spatial half of a
- /// depthwise-separable convolution. Follow with a 1×1 [`Conv2d`] to mix
- /// channels back together.
+ /// One filter per channel (`groups == channels`).
pub fn depthwise(channels: usize, kernel: usize, stride: usize, padding: usize) -> Conv2d {
Conv2d::grouped(channels, channels, kernel, stride, padding, channels)
}
@@ -111,10 +101,7 @@ impl Conv2d {
&[out_channels, in_channels / groups, kh, kw],
fan_in,
)),
- bias: bias.then(|| {
- let bound = 1.0 / (fan_in as f32).sqrt();
- Param::new(Tensor::uniform(&[out_channels], -bound, bound))
- }),
+ bias: bias.then(|| Param::new(uniform_bias(out_channels, fan_in))),
window,
in_channels,
out_channels,
@@ -148,9 +135,7 @@ impl Module for Conv2d {
let (out_h, out_w) = self.window.output_size(height, width);
let out_c = self.out_channels as i64;
- // [N, C·kh·kw, out_h·out_w] against [out_channels, C·kh·kw]; with
- // groups, the same product runs once per channel slice and the results
- // stack back along the channel axis.
+ // with groups the same product runs per channel slice and the results are stacked
let mut out = if self.groups == 1 {
let columns = input.im2col(self.window);
self.weight.tensor().reshape(&[out_c, -1]).matmul(&columns)
diff --git a/src/nn/conv1d.rs b/src/nn/conv1d.rs
index bfc1f6c..d22f8be 100644
--- a/src/nn/conv1d.rs
+++ b/src/nn/conv1d.rs
@@ -1,12 +1,4 @@
-//! 1-D convolution over sequences.
-//!
-//! A sequence is an image one pixel tall, so `Conv1d` *is* [`Conv2d`] with a
-//! `1×k` kernel — same unfold, same GEMM, same kernels on both devices. What
-//! this file adds is the sequence-specific option: causal padding, where all
-//! `(k−1)·d` pad positions go on the *left*, so output `t` reads only inputs
-//! `≤ t`. Stack causal layers with growing dilation and the receptive field
-//! doubles per layer while each stays one convolution — the WaveNet
-//! construction, and the convolutional counterpart of a causal mask.
+//! 1-D convolution, a [`Conv2d`] with a `1×k` kernel plus optional causal padding.
use crate::tensor::{Tensor, Window};
@@ -17,12 +9,12 @@ use super::param::Param;
/// Convolution over `[N, C, L]`.
pub struct Conv1d {
inner: Conv2d,
- /// Zeros prepended to the sequence; nonzero only for causal layers.
+ /// Zeros prepended to the sequence, nonzero only for causal layers.
left_pad: usize,
}
impl Conv1d {
- /// Symmetric padding, as in the 2-D case.
+ /// Symmetric padding.
pub fn new(
in_channels: usize,
out_channels: usize,
@@ -42,8 +34,7 @@ impl Conv1d {
}
}
- /// Causal, length-preserving, optionally dilated: output `t` depends only
- /// on inputs `≤ t`.
+ /// Causal and length-preserving: output `t` depends only on inputs `≤ t`.
pub fn causal(
in_channels: usize,
out_channels: usize,
@@ -53,7 +44,7 @@ impl Conv1d {
let window = Window {
kernel: (1, kernel),
stride: (1, 1),
- padding: (0, 0), // all padding is explicit, on the left
+ padding: (0, 0),
dilation: (1, dilation),
};
Conv1d {
@@ -72,8 +63,7 @@ impl Module for Conv1d {
input.shape()
);
- // Left-only zero padding, done explicitly: Window padding is symmetric
- // by design, and causality is exactly the asymmetric case.
+ // Window padding is symmetric, so the causal left pad is done by hand
let padded = if self.left_pad > 0 {
let zeros =
Tensor::zeros(&[input.dim(0), input.dim(1), self.left_pad]).to(input.device());
@@ -82,7 +72,7 @@ impl Module for Conv1d {
input.clone()
};
- let lifted = padded.unsqueeze(2); // [N, C, 1, L]
+ let lifted = padded.unsqueeze(2);
self.inner.forward(&lifted).squeeze(2)
}
diff --git a/src/nn/conv_transpose.rs b/src/nn/conv_transpose.rs
index 1d150cf..7348421 100644
--- a/src/nn/conv_transpose.rs
+++ b/src/nn/conv_transpose.rs
@@ -1,27 +1,18 @@
-//! Transposed convolution: the layer that makes images bigger.
+//! Transposed convolution, the forward pass of which is the convolution's
+//! backward pass: multiply by the transposed weight, then `col2im`.
//!
-//! A convolution's forward pass is `weight · im2col(x)`; its gradient scatters
-//! back through `col2im`. Transposed convolution simply *runs that adjoint as
-//! the forward pass* — multiply by the transposed weight, then `col2im` — so
-//! upsampling in a decoder or GAN generator is literally the gradient of the
-//! downsampling it mirrors. Nothing here is new machinery: both halves already
-//! exist, on both devices, with their derivatives.
-//!
-//! Sizes invert accordingly: a convolution maps `H → (H + 2p − span)/s + 1`,
-//! so this maps `H → (H−1)·s − 2p + span (+ output_padding)`, where
-//! `span = (k−1)·d + 1`. `output_padding` disambiguates strided shapes — with
-//! `s = 2`, inputs of 5 and 6 both convolve to 3, and it picks which one to
-//! return to.
+//! Output size is `(H−1)·s − 2p + (k−1)·d + 1 + output_padding`. The padding
+//! picks between inputs that convolve to the same size (5 and 6 at stride 2).
use crate::tensor::{Tensor, Window};
+use super::linear::uniform_bias;
use super::module::Module;
use super::param::Param;
/// Transposed 2-D convolution over `[N, C, H, W]`.
pub struct ConvTranspose2d {
- /// `[in_channels, out_channels, kh, kw]` — the mirror of [`Conv2d`](super::Conv2d)'s
- /// layout, because this layer is the mirror of its data flow.
+ /// `[in_channels, out_channels, kh, kw]`, the reverse of [`Conv2d`](super::Conv2d)'s layout.
pub weight: Param,
pub bias: Option,
window: Window,
@@ -55,19 +46,13 @@ impl ConvTranspose2d {
bias: bool,
) -> ConvTranspose2d {
let (kh, kw) = window.kernel;
- // Fan-in from the producing side: each output pixel receives from
- // roughly (out? no —) `in_channels·k²/s²` inputs; the conventional
- // init uses the weight's own receptive size, as Conv2d does.
let fan_in = in_channels * kh * kw;
ConvTranspose2d {
weight: Param::new(Tensor::kaiming_uniform(
&[in_channels, out_channels, kh, kw],
fan_in,
)),
- bias: bias.then(|| {
- let bound = 1.0 / (fan_in as f32).sqrt();
- Param::new(Tensor::uniform(&[out_channels], -bound, bound))
- }),
+ bias: bias.then(|| Param::new(uniform_bias(out_channels, fan_in))),
window,
in_channels,
out_channels,
@@ -75,8 +60,7 @@ impl ConvTranspose2d {
}
}
- /// Extra rows/columns on the output's far edge, for hitting an exact target
- /// size under stride. Must be smaller than the stride.
+ /// Extra rows and columns on the far edge. Must be smaller than the stride.
pub fn output_padding(mut self, padding: usize) -> ConvTranspose2d {
assert!(
padding < self.window.stride.0 && padding < self.window.stride.1,
@@ -132,9 +116,7 @@ impl Module for ConvTranspose2d {
let (batch, height, width) = (input.dim(0), input.dim(2), input.dim(3));
let (out_h, out_w) = self.output_size(height, width);
- // weightᵀ · x: [out_c·kh·kw, in_c] × [N, in_c, H·W] → columns, then
- // fold the columns into the larger image. matmul_tn reads the weight
- // transposed in place, and broadcasts it across the batch.
+ // matmul_tn reads the weight transposed in place and broadcasts it over the batch
let flat = input.reshape(&[
batch as i64,
self.in_channels as i64,
@@ -159,8 +141,7 @@ impl Module for ConvTranspose2d {
}
}
-/// Transposed 1-D convolution over `[N, C, L]`: the 2-D layer on a 1-pixel-tall
-/// image, which is not a shortcut but the definition — the maths is identical.
+/// Transposed 1-D convolution over `[N, C, L]`, the 2-D layer on a 1-pixel-tall image.
pub struct ConvTranspose1d {
inner: ConvTranspose2d,
}
diff --git a/src/nn/decoder.rs b/src/nn/decoder.rs
index 0bcbb79..1e3c103 100644
--- a/src/nn/decoder.rs
+++ b/src/nn/decoder.rs
@@ -1,11 +1,5 @@
-//! Transformer decoder blocks and the full encoder-decoder Transformer.
-//!
-//! A decoder block extends the encoder block with one more sub-layer: after
-//! causal self-attention, each target position attends over the encoder's
-//! output (the *memory*). That is what lets translation, summarisation, and
-//! any other sequence-to-sequence model read its input while writing its
-//! output. Decoder-only models like GPT do not need this file —
-//! [`TransformerStack::causal`] is the whole architecture.
+//! Decoder blocks with cross-attention over encoder output (the memory), and the
+//! encoder-decoder [`Transformer`]. Decoder-only models just use [`TransformerStack::causal`].
use crate::tensor::Tensor;
@@ -17,8 +11,7 @@ use super::norm::LayerNorm;
use super::param::Param;
use super::transformer::{Activation, TransformerStack};
-/// One pre-norm decoder block: causal self-attention, cross-attention over the
-/// encoder memory, then a feed-forward network, each wrapped in a residual.
+/// Pre-norm decoder block: causal self-attention, cross-attention, feed-forward.
pub struct TransformerDecoderBlock {
pub self_attention: MultiHeadAttention,
pub cross_attention: MultiHeadAttention,
@@ -56,11 +49,8 @@ impl TransformerDecoderBlock {
self.decode_masked(target, memory, None)
}
- /// [`decode`](Self::decode) with a padding mask for the memory:
- /// `[batch, source_len]`, 1 for real encoder positions and 0 for padding.
- ///
- /// The query side is normalized per sub-layer as usual; the memory arrives
- /// already normalized by the encoder's final norm and is used as-is.
+ /// [`decode`](Self::decode) with a memory padding mask, `[batch, source_len]`,
+ /// 1 for real positions. The memory is already normed by the encoder, so it is used as-is.
pub fn decode_masked(
&self,
target: &Tensor,
@@ -84,8 +74,7 @@ impl TransformerDecoderBlock {
}
impl Module for TransformerDecoderBlock {
- /// Degenerate single-input form: the target attends over itself as memory.
- /// Real sequence-to-sequence use goes through [`decode`](Self::decode).
+ /// Attends over the input as its own memory. Real use goes through [`decode`](Self::decode).
fn forward(&self, input: &Tensor) -> Tensor {
self.decode(input, input)
}
@@ -154,7 +143,7 @@ impl TransformerDecoder {
}
impl Module for TransformerDecoder {
- /// Degenerate single-input form; see [`TransformerDecoderBlock::forward`].
+ /// Attends over the input as its own memory.
fn forward(&self, input: &Tensor) -> Tensor {
self.decode(input, input)
}
@@ -177,12 +166,7 @@ impl Module for TransformerDecoder {
}
}
-/// The full encoder-decoder Transformer: a bidirectional encoder over the
-/// source, a causal decoder over the target with cross-attention into the
-/// encoder's output.
-///
-/// Embeddings and the output head stay outside, as they do for
-/// [`TransformerStack`] — this is the trunk, not the whole model.
+/// Encoder-decoder Transformer trunk. Embeddings and the output head stay outside.
pub struct Transformer {
pub encoder: TransformerStack,
pub decoder: TransformerDecoder,
@@ -214,9 +198,7 @@ impl Transformer {
self.run_masked(source, target, None)
}
- /// [`run`](Self::run) with a source padding mask, applied both to the
- /// encoder's self-attention and to the decoder's cross-attention — the two
- /// places a padded source position could leak in.
+ /// [`run`](Self::run) with a source padding mask, used by the encoder and the cross-attention.
pub fn run_masked(
&self,
source: &Tensor,
@@ -229,8 +211,7 @@ impl Transformer {
}
impl Module for Transformer {
- /// Degenerate single-input form: encode and decode the same sequence.
- /// Real sequence-to-sequence use goes through [`run`](Self::run).
+ /// Encodes and decodes the same sequence. Real use goes through [`run`](Self::run).
fn forward(&self, input: &Tensor) -> Tensor {
self.run(input, input)
}
diff --git a/src/nn/dropout.rs b/src/nn/dropout.rs
index c55858b..3795d5c 100644
--- a/src/nn/dropout.rs
+++ b/src/nn/dropout.rs
@@ -7,19 +7,15 @@ use crate::tensor::Tensor;
use super::module::Module;
-/// Zero each element with probability `p` during training, scaling the survivors
-/// by `1/(1-p)` so the expected activation is unchanged.
-///
-/// That rescaling is what lets eval mode be a no-op: the network sees the same
-/// average signal either way. Implemented as a multiply by a random mask, so the
-/// gradient is routed to exactly the units that survived.
+/// Zero each element with probability `p` while training, scaling survivors by
+/// `1/(1-p)` so eval mode can be a no-op.
pub struct Dropout {
p: f32,
training: AtomicBool,
}
impl Dropout {
- /// Panics unless `p` is in `[0, 1)` — `p = 1` would zero everything.
+ /// Panics unless `p` is in `[0, 1)`.
pub fn new(p: f32) -> Dropout {
assert!(
(0.0..1.0).contains(&p),
diff --git a/src/nn/embedding.rs b/src/nn/embedding.rs
index 8a0ff20..b5dcde7 100644
--- a/src/nn/embedding.rs
+++ b/src/nn/embedding.rs
@@ -5,10 +5,11 @@ use crate::tensor::{ops::index::index_from_f32, Tensor};
use super::module::Module;
use super::param::Param;
-/// A learnable `[vocab, dim]` table, looked up by token id.
-///
-/// The lookup is [`index_select`](Tensor::index_select), so the gradient
-/// scatter-adds back into the rows that were used and leaves the rest untouched.
+const INIT_STD: f32 = 0.02;
+/// Base of the geometric wavelength progression in the sinusoidal encoding.
+const WAVELENGTH_BASE: f32 = 10000.0;
+
+/// A learnable `[vocab, dim]` table looked up by token id.
pub struct Embedding {
pub weight: Param,
vocab: usize,
@@ -17,10 +18,9 @@ pub struct Embedding {
impl Embedding {
pub fn new(vocab: usize, dim: usize) -> Embedding {
- // Small normal init: large embeddings would dominate the residual stream
- // before the model has learned anything.
+ // small init so embeddings don't dominate the residual stream early on
Embedding {
- weight: Param::new(Tensor::randn(&[vocab, dim]).mul_scalar(0.02)),
+ weight: Param::new(Tensor::randn(&[vocab, dim]).mul_scalar(INIT_STD)),
vocab,
dim,
}
@@ -54,9 +54,8 @@ impl Embedding {
}
impl Module for Embedding {
- /// Treats finite, non-negative integer values as ids.
- /// Panics if an id cannot be represented as `usize`.
- /// Prefer [`lookup`](Embedding::lookup), which takes `usize` directly.
+ /// Treats finite, non-negative integer values as ids and panics on anything
+ /// else. Prefer [`lookup`](Embedding::lookup).
fn forward(&self, input: &Tensor) -> Tensor {
let ids: Vec = input.to_vec().into_iter().map(index_from_f32).collect();
let mut shape: Vec = input.shape().iter().map(|&d| d as i64).collect();
@@ -70,9 +69,6 @@ impl Module for Embedding {
}
/// Fixed sinusoidal position signal, added to token embeddings.
-///
-/// Not learned: the sin/cos pattern lets the model read relative offsets as a
-/// linear function of position, and generalises past the longest training length.
pub struct PositionalEncoding {
table: Tensor,
dim: usize,
@@ -83,7 +79,7 @@ impl PositionalEncoding {
let mut table = vec![0.0f32; max_len * dim];
for position in 0..max_len {
for i in (0..dim).step_by(2) {
- let angle = position as f32 / 10000f32.powf(i as f32 / dim as f32);
+ let angle = position as f32 / WAVELENGTH_BASE.powf(i as f32 / dim as f32);
table[position * dim + i] = angle.sin();
if i + 1 < dim {
table[position * dim + i + 1] = angle.cos();
diff --git a/src/nn/linear.rs b/src/nn/linear.rs
index 8cf8369..9238edf 100644
--- a/src/nn/linear.rs
+++ b/src/nn/linear.rs
@@ -5,11 +5,7 @@ use crate::tensor::Tensor;
use super::module::Module;
use super::param::Param;
-/// `y = x · Wᵀ + b`.
-///
-/// Accepts any input whose last dimension is `in_features`; leading dimensions
-/// are flattened into the batch and restored on the way out, so the same layer
-/// serves `[batch, features]` and `[batch, sequence, features]`.
+/// `y = x · Wᵀ + b` over the last dimension, any number of leading dimensions.
pub struct Linear {
pub weight: Param,
pub bias: Option,
@@ -31,8 +27,7 @@ impl Linear {
}
}
- /// A layer with no bias — the usual choice when a normalization layer with
- /// its own shift follows immediately.
+ /// A layer without bias.
pub fn no_bias(in_features: usize, out_features: usize) -> Linear {
Linear {
weight: Param::new(Tensor::kaiming_uniform(
@@ -64,7 +59,7 @@ impl Module for Linear {
input.shape()
);
- // `matmul_nt` consumes the weight as stored, [out, in], with no transposed copy.
+ // matmul_nt takes the weight as stored, [out, in], so no transposed copy
let flat = input.reshape(&[-1, self.in_features as i64]);
let mut out = flat.matmul_nt(&self.weight.tensor());
@@ -89,8 +84,8 @@ impl Module for Linear {
}
}
-/// PyTorch's linear bias init: uniform over `±1/√fan_in`.
-fn uniform_bias(out_features: usize, fan_in: usize) -> Tensor {
+/// Uniform over `±1/√fan_in`, the bias init shared with the conv layers.
+pub(super) fn uniform_bias(out_features: usize, fan_in: usize) -> Tensor {
let bound = 1.0 / (fan_in as f32).sqrt();
Tensor::uniform(&[out_features], -bound, bound)
}
diff --git a/src/nn/loss.rs b/src/nn/loss.rs
index 92958e7..1e9ca18 100644
--- a/src/nn/loss.rs
+++ b/src/nn/loss.rs
@@ -1,16 +1,12 @@
-//! Loss functions.
-//!
-//! Free functions, not layers — a loss has no parameters and no state, so
-//! `cross_entropy(&logits, &targets)` says everything a struct would.
+//! Loss functions, as free functions since they hold no state.
use crate::autograd::ops::loss::{CrossEntropyBackward, WeightedCrossEntropyBackward};
use crate::tensor::Tensor;
-/// Mean cross-entropy between `logits` `[batch, classes]` and integer class targets.
-///
-/// Takes raw logits, not probabilities: softmax and the log are fused here, so
-/// the numerically dangerous `log(exp(...))` never appears and a confidently
-/// wrong prediction yields a large finite loss instead of infinity.
+/// Keeps probabilities off 0 and 1 so `log` stays finite.
+const CLAMP_EPS: f32 = 1e-7;
+
+/// Mean cross-entropy of raw `logits` `[batch, classes]` against integer class targets.
pub fn cross_entropy(logits: &Tensor, targets: &[usize]) -> Tensor {
assert_eq!(
logits.ndim(),
@@ -38,7 +34,7 @@ pub fn cross_entropy(logits: &Tensor, targets: &[usize]) -> Tensor {
let base = row * classes;
let values = &data[base..base + classes];
- // Shift by the row maximum so exp() cannot overflow.
+ // shift by the row max so exp can't overflow
let shift = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
let sum: f32 = values.iter().map(|&v| (v - shift).exp()).sum();
for c in 0..classes {
@@ -63,19 +59,14 @@ pub fn mse(prediction: &Tensor, target: &Tensor) -> Tensor {
prediction.sub(target).square().mean()
}
-/// Mean absolute error. Less sensitive to outliers than [`mse`].
+/// Mean absolute error.
pub fn mae(prediction: &Tensor, target: &Tensor) -> Tensor {
prediction.sub(target).abs().mean()
}
-/// Binary cross-entropy over probabilities already in `[0, 1]`.
-///
-/// Predictions are clamped away from the endpoints, because `log(0)` is `-inf`
-/// and one saturated output would poison the whole batch. Prefer
-/// [`bce_with_logits`] when you have raw scores.
+/// Binary cross-entropy over probabilities. Prefer [`bce_with_logits`] for raw scores.
pub fn bce(prediction: &Tensor, target: &Tensor) -> Tensor {
- const EDGE: f32 = 1e-7;
- let p = prediction.clamp(EDGE, 1.0 - EDGE);
+ let p = prediction.clamp(CLAMP_EPS, 1.0 - CLAMP_EPS);
let positive = target.mul(&p.log());
let negative = target
.neg()
@@ -84,23 +75,14 @@ pub fn bce(prediction: &Tensor, target: &Tensor) -> Tensor {
positive.add(&negative).neg().mean()
}
-/// Binary cross-entropy straight from logits.
-///
-/// Uses `max(x,0) - x·t + log(1 + e^-|x|)`, which is the same value as
-/// `bce(sigmoid(x), t)` but never exponentiates a large positive number.
+/// Binary cross-entropy from logits, as `max(x,0) - x·t + log(1 + e^-|x|)` to avoid overflow.
pub fn bce_with_logits(logits: &Tensor, target: &Tensor) -> Tensor {
let floor = logits.clamp(0.0, f32::INFINITY);
let stable_log = logits.abs().neg().exp().add_scalar(1.0).log();
floor.sub(&logits.mul(target)).add(&stable_log).mean()
}
-// ── Regression ───────────────────────────────────────────────────────────────
-
/// Huber loss: quadratic within `delta` of the target, linear beyond it.
-///
-/// Written branch-free as `½·min(|e|, δ)² + δ·(|e| − min(|e|, δ))`, which
-/// equals `½e²` inside the threshold and `δ(|e| − ½δ)` outside — MSE's smooth
-/// gradient near the answer, MAE's bounded gradient for outliers.
pub fn huber(prediction: &Tensor, target: &Tensor, delta: f32) -> Tensor {
assert!(delta > 0.0, "huber delta must be positive, got {delta}");
let residual = prediction.sub(target).abs();
@@ -110,17 +92,13 @@ pub fn huber(prediction: &Tensor, target: &Tensor, delta: f32) -> Tensor {
quadratic.add(&linear).mean()
}
-/// Huber with `delta = 1`, under the name detection literature uses.
+/// Huber with `delta = 1`.
pub fn smooth_l1(prediction: &Tensor, target: &Tensor) -> Tensor {
huber(prediction, target, 1.0)
}
-/// Negative log-likelihood of a Gaussian with predicted mean and log-variance,
-/// per element: `½(log σ² + (x − μ)²/σ²)`, dropping the `½ln 2π` constant.
-///
-/// Predicting `log σ²` rather than `σ²` keeps the variance positive by
-/// construction and the loss finite everywhere — the standard trick for
-/// heteroscedastic regression.
+/// Gaussian negative log-likelihood from a predicted mean and log-variance,
+/// without the `½ln 2π` constant.
pub fn gaussian_nll(mean: &Tensor, log_variance: &Tensor, target: &Tensor) -> Tensor {
let squared_error = target.sub(mean).square();
let precision = log_variance.neg().exp();
@@ -130,32 +108,21 @@ pub fn gaussian_nll(mean: &Tensor, log_variance: &Tensor, target: &Tensor) -> Te
.mean()
}
-/// Negative log-likelihood of a Poisson with predicted log-rate, per element:
-/// `e^{log λ} − k·log λ`, dropping the `log k!` term that no gradient reaches.
+/// Poisson negative log-likelihood from a predicted log-rate, without the `log k!` term.
pub fn poisson_nll(log_rate: &Tensor, target: &Tensor) -> Tensor {
log_rate.exp().sub(&target.mul(log_rate)).mean()
}
-// ── Distributions ────────────────────────────────────────────────────────────
-
-/// `KL(target ‖ prediction)`: rows of probabilities, summed over the last
-/// axis and averaged over the rest.
-///
-/// Both sides are clamped away from zero so `0·log 0` — which is 0 in the
-/// limit — cannot become `NaN` in f32. The direction matters: this penalises
-/// the prediction for missing mass wherever the *target* has some.
+/// `KL(target ‖ prediction)` over rows of probabilities, summed on the last axis.
+/// Both sides are clamped so `0·log 0` can't become `NaN`.
pub fn kl_divergence(prediction: &Tensor, target: &Tensor) -> Tensor {
- const EDGE: f32 = 1e-7;
- let p = prediction.clamp(EDGE, 1.0);
- let t = target.clamp(EDGE, 1.0);
+ let p = prediction.clamp(CLAMP_EPS, 1.0);
+ let t = target.clamp(CLAMP_EPS, 1.0);
let rows = (p.numel() / p.last_dim()) as f32;
t.mul(&t.log().sub(&p.log())).sum().div_scalar(rows)
}
-/// Negative log-likelihood over rows of *log*-probabilities, the pairing for a
-/// model that ends in `log_softmax`. `nll(log_softmax(x), t)` equals
-/// `cross_entropy(x, t)`; use the fused form unless the log-probabilities are
-/// needed elsewhere too.
+/// Negative log-likelihood over rows of log-probabilities. Equals [`cross_entropy`] after `log_softmax`.
pub fn nll(log_probs: &Tensor, targets: &[usize]) -> Tensor {
let (batch, classes) = (log_probs.dim(0), log_probs.dim(1));
assert_eq!(
@@ -165,8 +132,7 @@ pub fn nll(log_probs: &Tensor, targets: &[usize]) -> Tensor {
targets.len()
);
- // Selecting one entry per row is a dot with the one-hot targets — a
- // constant, so the graph reaches only the log-probabilities.
+ // picking one entry per row as a dot with constant one-hot targets
let mut one_hot = vec![0.0f32; batch * classes];
for (row, &target) in targets.iter().enumerate() {
assert!(
@@ -183,18 +149,10 @@ pub fn nll(log_probs: &Tensor, targets: &[usize]) -> Tensor {
.neg()
}
-// ── Classification under imbalance ───────────────────────────────────────────
-
-/// Binary focal loss from logits (Lin et al., 2017).
-///
-/// Scales each example's BCE by `(1 − p_correct)^γ`: an example the model
-/// already gets right contributes almost nothing, so the rare hard positives
-/// dominate the gradient instead of drowning under easy negatives. `alpha`
-/// additionally reweights the positive class; `gamma = 2`, `alpha = 0.25` are
-/// the paper's defaults.
+/// Binary focal loss from logits: BCE scaled by `(1 − p_correct)^gamma`, with
+/// `alpha` weighting the positive class. The paper uses `gamma = 2`, `alpha = 0.25`.
pub fn focal_bce_with_logits(logits: &Tensor, target: &Tensor, gamma: f32, alpha: f32) -> Tensor {
- const EDGE: f32 = 1e-7;
- let p = logits.sigmoid().clamp(EDGE, 1.0 - EDGE);
+ let p = logits.sigmoid().clamp(CLAMP_EPS, 1.0 - CLAMP_EPS);
let one_minus_p = p.neg().add_scalar(1.0);
let one_minus_t = target.neg().add_scalar(1.0);
@@ -210,10 +168,6 @@ pub fn focal_bce_with_logits(logits: &Tensor, target: &Tensor, gamma: f32, alpha
}
/// Dice loss over probabilities: `1 − 2|P∩T| / (|P| + |T|)`, smoothed.
-///
-/// Measures overlap rather than per-pixel agreement, so a tiny foreground
-/// region weighs as much as a large one — why segmentation reaches for it
-/// when cross-entropy collapses to "predict background everywhere".
pub fn dice(prediction: &Tensor, target: &Tensor) -> Tensor {
const SMOOTH: f32 = 1.0;
let intersection = prediction
@@ -225,20 +179,18 @@ pub fn dice(prediction: &Tensor, target: &Tensor) -> Tensor {
intersection.div(&total).neg().add_scalar(1.0)
}
-// ── Configurable cross-entropy ───────────────────────────────────────────────
-
/// How per-sample losses become the reported value.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Reduction {
- /// Weighted mean — the default, and what a learning rate is tuned against.
+ /// Weighted mean.
Mean,
/// Plain sum; scales with batch size.
Sum,
- /// No reduction: one loss per row, for masking or reweighting downstream.
+ /// One loss per row.
None,
}
-/// [`cross_entropy`] with the knobs real datasets need.
+/// [`cross_entropy`] with label smoothing, class weights and an ignored index.
///
/// ```
/// # use fastnn::prelude::*;
@@ -249,12 +201,8 @@ pub enum Reduction {
/// let loss = loss_fn.compute(&logits, &[2, 0, 4]);
/// ```
///
-/// With label smoothing `s`, the target distribution becomes
-/// `(1−s)·one_hot + s/K` — the model is asked for confidence `1−s+s/K`, not
-/// certainty, which keeps logits bounded and calibration honest. Class weights
-/// follow the weighted-mean convention: the loss is divided by the sum of the
-/// participating rows' weights, so reweighting does not silently rescale the
-/// learning rate. Ignored rows contribute nothing to loss or gradient.
+/// Under `Mean` the loss is divided by the summed weights of the rows that
+/// count, so weighting doesn't rescale the learning rate.
pub struct CrossEntropyLoss {
label_smoothing: f32,
class_weights: Option>,
@@ -272,7 +220,7 @@ impl CrossEntropyLoss {
}
}
- /// Blend `s` of the probability mass uniformly across classes.
+ /// Spread `s` of the probability mass uniformly over the classes.
pub fn label_smoothing(mut self, s: f32) -> CrossEntropyLoss {
assert!(
(0.0..1.0).contains(&s),
@@ -282,13 +230,13 @@ impl CrossEntropyLoss {
self
}
- /// Per-class weights, e.g. inverse class frequency for imbalanced data.
+ /// Per-class weights.
pub fn class_weights(mut self, weights: Vec) -> CrossEntropyLoss {
self.class_weights = Some(weights);
self
}
- /// Rows with this target contribute nothing — padding tokens, typically.
+ /// Rows with this target contribute nothing.
pub fn ignore_index(mut self, index: usize) -> CrossEntropyLoss {
self.ignore_index = Some(index);
self
@@ -299,8 +247,7 @@ impl CrossEntropyLoss {
self
}
- /// `[batch, classes]` logits and one target per row → the reduced loss
- /// (or `[batch]` under [`Reduction::None`]).
+ /// Loss for `[batch, classes]` logits, or `[batch]` under [`Reduction::None`].
pub fn compute(&self, logits: &Tensor, targets: &[usize]) -> Tensor {
assert_eq!(
logits.ndim(),
@@ -342,8 +289,7 @@ impl CrossEntropyLoss {
weight_total += weight;
let values = &data[row * classes..(row + 1) * classes];
- // log p_c = v_c − shift − ln Σ e^{v−shift}: the shift makes every
- // exponent ≤ 0, so nothing overflows however large the logits.
+ // shift by the row max so every exponent is <= 0
let shift = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
let log_sum = values.iter().map(|&v| (v - shift).exp()).sum::().ln();
@@ -352,7 +298,7 @@ impl CrossEntropyLoss {
let log_prob = value - shift - log_sum;
let q = if class == target { on_target } else { uniform };
loss -= q * log_prob;
- // d loss / d v_c = p_c − q_c, scaled by this row's weight.
+ // d loss / d logit is p - q, scaled by the row weight
difference[row * classes + class] = weight * ((log_prob).exp() - q);
}
per_row[row] = weight * loss;
@@ -361,8 +307,7 @@ impl CrossEntropyLoss {
let difference = Tensor::from_vec(difference, &[batch, classes]);
let (value, normalizer, per_row_reduction) = match self.reduction {
Reduction::Mean => {
- // An all-ignored batch has nothing to average; 0 with a zero
- // gradient is the only answer that does not poison the run.
+ // all rows ignored: report 0 with zero gradient instead of dividing by 0
let denominator = if weight_total > 0.0 {
weight_total
} else {
diff --git a/src/nn/metric_losses.rs b/src/nn/metric_losses.rs
index 7ba35a2..4c9f9c7 100644
--- a/src/nn/metric_losses.rs
+++ b/src/nn/metric_losses.rs
@@ -1,29 +1,18 @@
-//! Losses over embedding geometry: pull related points together, push
-//! unrelated ones apart.
-//!
-//! Classification losses grade a prediction against a label; these grade the
-//! *shape of the embedding space* itself — distances and angles between pairs
-//! and triples — which is what retrieval, verification, and self-supervised
-//! pretraining actually optimize. All of them compose from ordinary
-//! differentiable ops, so their gradients come from the existing rules.
+//! Losses on embedding geometry (pairs and triples), built from ordinary ops.
use crate::tensor::Tensor;
use super::loss::cross_entropy;
-/// Guards divisions and square roots at zero distance, where the true
-/// derivative of ‖·‖₂ does not exist.
+/// The derivative of a norm is undefined at zero, so keep sqrt and division off it.
const EPS: f32 = 1e-8;
-/// Row-wise Euclidean distance between `[batch, dim]` tensors → `[batch, 1]`.
-///
-/// `√(Σd² + ε)` rather than `√Σd²`: at zero distance the exact root has an
-/// undefined gradient, and the ε turns it into a smooth minimum instead.
+/// Row-wise Euclidean distance, `[batch, dim]` to `[batch, 1]`.
fn row_distance(a: &Tensor, b: &Tensor) -> Tensor {
a.sub(b).square().sum_axis_keep(1).add_scalar(EPS).sqrt()
}
-/// Row-wise cosine similarity between `[batch, dim]` tensors → `[batch, 1]`.
+/// Row-wise cosine similarity, `[batch, dim]` to `[batch, 1]`.
fn row_cosine(a: &Tensor, b: &Tensor) -> Tensor {
let dot = a.mul(b).sum_axis_keep(1);
let norms = a
@@ -43,12 +32,8 @@ fn mask(targets: &[f32], predicate: impl Fn(f32) -> bool, device: crate::tensor:
Tensor::from_vec(values, &[targets.len(), 1]).to(device)
}
-/// Cosine embedding loss: for each row pair, `1 − cos` when `target = +1`,
-/// `max(0, cos − margin)` when `target = −1`.
-///
-/// Similar pairs are pulled to angle zero; dissimilar ones are only pushed
-/// until their similarity drops below the margin — beyond that they stop
-/// mattering, which is what keeps the space from collapsing outward forever.
+/// Cosine embedding loss: `1 − cos` when `target = +1`, `max(0, cos − margin)`
+/// when `target = −1`.
pub fn cosine_embedding(a: &Tensor, b: &Tensor, targets: &[f32], margin: f32) -> Tensor {
assert_eq!(
a.dim(0),
@@ -68,10 +53,6 @@ pub fn cosine_embedding(a: &Tensor, b: &Tensor, targets: &[f32], margin: f32) ->
}
/// Triplet margin loss: `max(0, d(anchor, positive) − d(anchor, negative) + margin)`.
-///
-/// The hinge goes silent once the negative is `margin` farther than the
-/// positive — only violating triplets produce gradient, so what you mine into
-/// the batch is what the model learns from.
pub fn triplet_margin(
anchor: &Tensor,
positive: &Tensor,
@@ -87,8 +68,7 @@ pub fn triplet_margin(
.mean()
}
-/// Margin ranking loss: `max(0, −t·(x₁ − x₂) + margin)` with `t = ±1` saying
-/// which input should score higher. The SVM hinge, applied to rankings.
+/// Margin ranking loss: `max(0, −t·(x₁ − x₂) + margin)`, `t = ±1` picks the higher input.
pub fn margin_ranking(x1: &Tensor, x2: &Tensor, targets: &[f32], margin: f32) -> Tensor {
assert_eq!(
x1.numel(),
@@ -104,11 +84,7 @@ pub fn margin_ranking(x1: &Tensor, x2: &Tensor, targets: &[f32], margin: f32) ->
difference.mul(&sign).neg().add_scalar(margin).relu().mean()
}
-/// Contrastive pair loss (Hadsell et al., 2006): `t·d² + (1−t)·max(0, m − d)²`
-/// with `t = 1` for pairs that belong together.
-///
-/// Squared distance pulls matching pairs together; non-matching pairs are
-/// pushed quadratically until they clear the margin `m`, then released.
+/// Contrastive pair loss: `t·d² + (1−t)·max(0, m − d)²`, `t = 1` for matching pairs.
pub fn contrastive(a: &Tensor, b: &Tensor, targets: &[f32], margin: f32) -> Tensor {
assert_eq!(
a.dim(0),
@@ -132,13 +108,7 @@ pub fn contrastive(a: &Tensor, b: &Tensor, targets: &[f32], margin: f32) -> Tens
pull.add(&push).mean()
}
-/// InfoNCE, the loss behind CLIP and SimCLR-style pretraining.
-///
-/// Row `i` of `queries` and row `i` of `keys` are the matching pair. Both sets
-/// are L2-normalized, all pairwise similarities become logits at temperature
-/// `τ`, and each query must classify its own key out of the whole batch — so
-/// every other row serves as a free negative, and the effective task gets
-/// harder as the batch grows. Low `τ` sharpens the contrast; 0.07 is typical.
+/// InfoNCE. Row `i` of `queries` matches row `i` of `keys`, every other row is a negative.
pub fn info_nce(queries: &Tensor, keys: &Tensor, temperature: f32) -> Tensor {
assert!(
temperature > 0.0,
@@ -151,7 +121,7 @@ pub fn info_nce(queries: &Tensor, keys: &Tensor, temperature: f32) -> Tensor {
.matmul_nt(&normalize(keys))
.div_scalar(temperature);
- // The matching key for query i sits at column i: the targets are the diagonal.
+ // query i's match is column i
let diagonal: Vec = (0..batch).collect();
cross_entropy(&logits, &diagonal)
}
diff --git a/src/nn/mod.rs b/src/nn/mod.rs
index f15c9a0..fb259ae 100644
--- a/src/nn/mod.rs
+++ b/src/nn/mod.rs
@@ -1,32 +1,5 @@
-//! Neural network layers.
-//!
-//! Everything implements [`Module`]. A layer's weights are [`Param`] handles —
-//! shared slots, not owned values — which is why the optimizer can hold them at
-//! the same time as the model:
-//!
-//! ```
-//! use fastnn::prelude::*;
-//!
-//! let model = Sequential::new()
-//! .add(Linear::new(8, 16))
-//! .add(ReLU)
-//! .add(Linear::new(16, 2));
-//!
-//! let mut opt = SGD::new(model.parameters(), 0.1);
-//!
-//! let logits = model.forward(&Tensor::randn(&[4, 8]));
-//! let loss = cross_entropy(&logits, &[0, 1, 1, 0]);
-//!
-//! opt.zero_grad();
-//! loss.backward();
-//! opt.step();
-//! ```
-//!
-//! ## Writing a layer
-//!
-//! Implement [`forward`](Module::forward), and [`named_parameters`](Module::named_parameters)
-//! if it has weights. Device placement, parameter counting, and gradient zeroing
-//! come from those for free.
+//! Layers, losses and decoding helpers. Everything implements [`Module`], and a
+//! layer's weights are shared [`Param`] handles so an optimizer can hold them too.
pub mod activation;
pub mod attention;
diff --git a/src/nn/module.rs b/src/nn/module.rs
index 44505a3..a6814b7 100644
--- a/src/nn/module.rs
+++ b/src/nn/module.rs
@@ -4,15 +4,10 @@ use crate::tensor::{Device, Tensor};
use super::param::{Buffer, Param};
-/// A piece of a model.
+/// A piece of a model. Only [`forward`](Module::forward) is required, plus
+/// [`named_parameters`](Module::named_parameters) if it has weights.
///
-/// Only [`forward`](Module::forward) is required. A layer with weights also
-/// implements [`named_parameters`](Module::named_parameters); everything else —
-/// parameter counting, device placement, gradient zeroing — falls out of that.
-///
-/// Every method takes `&self`. Parameters carry their own interior mutability,
-/// so a container can recurse into its children without threading `&mut` through
-/// the whole tree, and `model.eval()` needs no `mut` binding.
+/// Everything takes `&self`, since parameters have interior mutability.
///
/// ```
/// use fastnn::prelude::*;
@@ -32,9 +27,7 @@ pub trait Module: Send + Sync {
/// Transform an input into an output.
fn forward(&self, input: &Tensor) -> Tensor;
- /// This module's learnable tensors, with names like `"weight"` or `"2.bias"`.
- ///
- /// A container prefixes its children's names; a leaf layer returns its own.
+ /// Learnable tensors, named like `"weight"` or `"2.bias"`.
fn named_parameters(&self) -> Vec<(String, Param)> {
Vec::new()
}
@@ -44,15 +37,11 @@ pub trait Module: Send + Sync {
Vec::new()
}
- /// Switch between training and inference behaviour.
- ///
- /// Only [`Dropout`](super::Dropout) and [`BatchNorm2d`](super::BatchNorm2d)
- /// act on this; containers pass it down.
+ /// Switch between training and inference behaviour. Only
+ /// [`Dropout`](super::Dropout) and [`BatchNorm2d`](super::BatchNorm2d) act on it.
fn set_training(&self, _training: bool) {}
- // ── Provided ─────────────────────────────────────────────────────────────
-
- /// Every learnable tensor. Hand this straight to an optimizer.
+ /// Every learnable tensor, ready to hand to an optimizer.
fn parameters(&self) -> Vec {
self.named_parameters()
.into_iter()
@@ -75,17 +64,14 @@ pub trait Module: Send + Sync {
}
}
- /// Stop training every parameter in this module.
- ///
- /// Freeze a pretrained backbone, leave the head trainable, and the optimizer
- /// updates only the head — the backward pass does not even compute the rest.
+ /// Freeze every parameter.
fn freeze(&self) {
for (_, p) in self.named_parameters() {
p.freeze();
}
}
- /// Resume training every parameter in this module.
+ /// Unfreeze every parameter.
fn unfreeze(&self) {
for (_, p) in self.named_parameters() {
p.unfreeze();
@@ -107,18 +93,18 @@ pub trait Module: Send + Sync {
}
}
- /// Enable training behaviour (dropout active, batch-norm using batch statistics).
+ /// Training mode.
fn train(&self) {
self.set_training(true);
}
- /// Enable inference behaviour.
+ /// Inference mode.
fn eval(&self) {
self.set_training(false);
}
}
-/// Prefix each name with `scope`, for a container building parameter paths.
+/// Prefix each name with `scope`.
///
/// ```
/// # use fastnn::prelude::*;
diff --git a/src/nn/norm.rs b/src/nn/norm.rs
index 1205e04..f9c9928 100644
--- a/src/nn/norm.rs
+++ b/src/nn/norm.rs
@@ -7,10 +7,12 @@ use crate::tensor::Tensor;
use super::module::Module;
use super::param::{Buffer, Param};
+const LAYER_NORM_EPS: f32 = 1e-5;
+const RMS_NORM_EPS: f32 = 1e-6;
+const BATCH_NORM_EPS: f32 = 1e-5;
+const BATCH_NORM_MOMENTUM: f32 = 0.1;
+
/// Normalize each sample over its last dimension, then scale and shift.
-///
-/// The workhorse of transformers: it depends on nothing but the sample itself,
-/// so batch size and sequence length cannot change its behaviour.
pub struct LayerNorm {
pub gamma: Param,
pub beta: Param,
@@ -24,11 +26,11 @@ impl LayerNorm {
gamma: Param::new(Tensor::ones(&[size])),
beta: Param::new(Tensor::zeros(&[size])),
size,
- eps: 1e-5,
+ eps: LAYER_NORM_EPS,
}
}
- /// Override the variance floor. Larger values trade accuracy for stability.
+ /// Override the variance floor.
pub fn eps(mut self, eps: f32) -> LayerNorm {
self.eps = eps;
self
@@ -56,9 +58,6 @@ impl Module for LayerNorm {
}
/// Scale by the root-mean-square of each sample, with no mean subtraction.
-///
-/// Cheaper than [`LayerNorm`] and used by most recent large models. Built from
-/// ordinary tensor ops, so it needs no gradient rule of its own.
pub struct RMSNorm {
pub gamma: Param,
size: usize,
@@ -70,7 +69,7 @@ impl RMSNorm {
RMSNorm {
gamma: Param::new(Tensor::ones(&[size])),
size,
- eps: 1e-6,
+ eps: RMS_NORM_EPS,
}
}
}
@@ -98,11 +97,8 @@ impl Module for RMSNorm {
}
}
-/// Normalize each channel of `[N, C, H, W]` across the batch and spatial positions.
-///
-/// In training it uses the current batch's statistics and folds them into a
-/// running estimate; in eval it applies that estimate as a plain affine map, so
-/// a single image gives the same answer as a large batch.
+/// Normalize each channel of `[N, C, H, W]` over the batch and spatial positions.
+/// Eval mode uses the running statistics instead of the batch's.
pub struct BatchNorm2d {
pub gamma: Param,
pub beta: Param,
@@ -122,8 +118,8 @@ impl BatchNorm2d {
running_mean: Buffer::new(Tensor::zeros(&[channels])),
running_var: Buffer::new(Tensor::ones(&[channels])),
channels,
- eps: 1e-5,
- momentum: 0.1,
+ eps: BATCH_NORM_EPS,
+ momentum: BATCH_NORM_MOMENTUM,
training: AtomicBool::new(true),
}
}
@@ -134,7 +130,6 @@ impl BatchNorm2d {
self
}
- /// Blend this batch's statistics into the running estimates.
fn update_running(&self, mean: &[f32], var: &[f32]) {
let blend = |old: Tensor, new: &[f32]| {
let old = old.to_vec();
@@ -165,7 +160,6 @@ impl Module for BatchNorm2d {
let per_channel = [1i64, self.channels as i64, 1, 1];
if !self.training.load(Ordering::Relaxed) {
- // Frozen statistics make this an affine map, which ordinary ops express.
let scale = self.running_var.value().add_scalar(self.eps).sqrt();
let centred = input.sub(&self.running_mean.value().reshape(&per_channel));
return centred
diff --git a/src/nn/param.rs b/src/nn/param.rs
index 82ff84b..7589052 100644
--- a/src/nn/param.rs
+++ b/src/nn/param.rs
@@ -1,18 +1,6 @@
-//! Shared handles to the tensors a model owns.
-//!
-//! A [`Param`] is a slot, not a value. Cloning one gives a second handle to the
-//! same weight and the same gradient, which is why `Adam::new(model.parameters())`
-//! works: the optimizer holds its own handles and updates them in place, with no
-//! borrow of the model at step time.
-//!
-//! ```text
-//! Linear { weight: Param ─┐
-//! ├─► Arc> the weight
-//! Adam { params: [Param]┘ Arc>> its gradient
-//! ```
-//!
-//! [`Buffer`] is the same sharing without a gradient: state that persists and is
-//! checkpointed but is not learned, such as batch-norm running statistics.
+//! Shared handles to a model's tensors. Cloning a [`Param`] gives another handle
+//! to the same weight and gradient, which is how an optimizer updates a model
+//! it does not borrow. [`Buffer`] is the same without a gradient.
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, RwLock};
@@ -39,12 +27,8 @@ impl Param {
}
}
- /// This parameter as a tensor, for use in a forward pass.
- ///
- /// Normally a graph leaf, so gradients reaching it land in this parameter's
- /// slot. While frozen it comes back detached instead — the backward pass then
- /// stops at this point on its own, costing nothing to compute a gradient that
- /// would only be discarded.
+ /// The parameter as a graph leaf for a forward pass. Detached while frozen,
+ /// so backward never computes a gradient for it.
pub fn tensor(&self) -> Tensor {
let value = self.value.read().unwrap();
if !self.is_trainable() {
@@ -58,17 +42,13 @@ impl Param {
)
}
- /// Stop training this parameter. Optimizers skip it and no gradient is
- /// computed for it at all.
- ///
- /// The usual reason is transfer learning: freeze a pretrained backbone and
- /// train only the new head on top.
+ /// Stop training this parameter.
pub fn freeze(&self) {
self.trainable.store(false, Ordering::Relaxed);
self.grad.clear();
}
- /// Resume training this parameter.
+ /// Undo [`freeze`](Param::freeze).
pub fn unfreeze(&self) {
self.trainable.store(true, Ordering::Relaxed);
}
@@ -83,7 +63,7 @@ impl Param {
self.value.read().unwrap().detach()
}
- /// Replace the value. Every handle sees the change; the gradient is untouched.
+ /// Replace the value for every handle. Shape must not change.
pub fn set_value(&self, value: Tensor) {
assert_eq!(
value.shape(),
@@ -103,15 +83,13 @@ impl Param {
self.grad.clear();
}
- /// Replace the gradient outright. Gradient clipping and custom optimizers
- /// use this; the backward pass itself accumulates instead.
+ /// Replace the gradient, for clipping and custom optimizers.
pub fn set_grad(&self, grad: Tensor) {
self.grad.clear();
self.grad.accumulate(grad.detach());
}
- /// Move the value onto `device` in place. The gradient is dropped, since a
- /// stale gradient on the old device could never be applied.
+ /// Move the value to `device`, dropping the gradient.
pub fn to_device(&self, device: Device) {
let mut value = self.value.write().unwrap();
if value.device() != device {
@@ -139,7 +117,7 @@ impl std::fmt::Debug for Param {
}
}
-/// Persistent state that is saved and restored but never learned.
+/// State that is checkpointed but not learned.
#[derive(Clone)]
pub struct Buffer {
value: Arc>,
diff --git a/src/nn/pooling.rs b/src/nn/pooling.rs
index 025bb49..1bcc69b 100644
--- a/src/nn/pooling.rs
+++ b/src/nn/pooling.rs
@@ -10,7 +10,7 @@ pub struct MaxPool2d {
}
impl MaxPool2d {
- /// A `kernel × kernel` window that does not overlap — the usual downsampler.
+ /// A non-overlapping `kernel × kernel` window.
pub fn new(kernel: usize) -> MaxPool2d {
MaxPool2d {
window: Window::square(kernel, kernel, 0),
diff --git a/src/nn/rnn.rs b/src/nn/rnn.rs
index 36b02ce..3426a27 100644
--- a/src/nn/rnn.rs
+++ b/src/nn/rnn.rs
@@ -1,20 +1,12 @@
-//! Recurrent layers.
-//!
-//! Both cells are written with ordinary differentiable tensor ops, so they train
-//! through backpropagation-through-time with no gradient code of their own. The
-//! cost is one graph node per gate per timestep: correct, and fine for short
-//! sequences, but a transformer will be far faster on long ones.
-//!
-//! Each layer is single-layer and unidirectional. Stack them in a
-//! [`Sequential`](super::Sequential) — the `[batch, sequence, hidden]` output
-//! feeds straight into the next.
+//! Single-layer, unidirectional LSTM and GRU built from ordinary tensor ops, so
+//! one graph node per gate per timestep. Stack them in a [`Sequential`](super::Sequential).
use crate::tensor::Tensor;
use super::module::Module;
use super::param::Param;
-/// The four weight tensors every gated recurrent cell has.
+/// Weights shared by both cells.
struct Gates {
input_weight: Param,
hidden_weight: Param,
@@ -56,10 +48,6 @@ impl Gates {
}
/// Long short-term memory.
-///
-/// Carries a cell state alongside the hidden state; the forget gate multiplies
-/// it rather than replacing it, which is what keeps gradients alive over long
-/// spans.
pub struct LSTM {
gates: Gates,
input_size: usize,
@@ -69,9 +57,7 @@ pub struct LSTM {
impl LSTM {
pub fn new(input_size: usize, hidden_size: usize) -> LSTM {
let gates = Gates::new(input_size, hidden_size, 4);
- // Start the forget gate open. At zero bias the sigmoid sits at 0.5 and
- // halves the cell state every step, so early gradients vanish before the
- // model has learned what to keep.
+ // forget gate starts open, at zero bias it halves the cell state every step
let mut bias = gates.input_bias.value().to_vec();
bias[hidden_size..2 * hidden_size].fill(1.0);
gates
@@ -85,10 +71,7 @@ impl LSTM {
}
}
- /// Run over `[batch, sequence, input_size]`.
- ///
- /// Returns the per-step outputs `[batch, sequence, hidden_size]` plus the
- /// final hidden and cell states.
+ /// Run over `[batch, sequence, input_size]`, returning outputs and the final `(h, c)`.
pub fn run(
&self,
input: &Tensor,
@@ -132,7 +115,7 @@ impl Module for LSTM {
}
}
-/// Gated recurrent unit — LSTM's gating with no separate cell state.
+/// Gated recurrent unit, LSTM's gating with no separate cell state.
pub struct GRU {
gates: Gates,
input_size: usize,
@@ -148,8 +131,7 @@ impl GRU {
}
}
- /// Run over `[batch, sequence, input_size]`, returning the outputs and the
- /// final hidden state.
+ /// Run over `[batch, sequence, input_size]`, returning outputs and the final hidden state.
pub fn run(&self, input: &Tensor, initial: Option<&Tensor>) -> (Tensor, Tensor) {
let (batch, steps) = check_sequence(input, self.input_size, "GRU");
let mut h = initial
@@ -166,8 +148,7 @@ impl GRU {
};
let reset = part(&from_input, 0).add(&part(&from_hidden, 0)).sigmoid();
let update = part(&from_input, 1).add(&part(&from_hidden, 1)).sigmoid();
- // The reset gate scales the *hidden* contribution only, letting the
- // candidate ignore history without ignoring the current input.
+ // reset scales the hidden contribution only
let candidate = part(&from_input, 2)
.add(&reset.mul(&part(&from_hidden, 2)))
.tanh();
diff --git a/src/nn/sample.rs b/src/nn/sample.rs
index 6fdd38b..cf1de96 100644
--- a/src/nn/sample.rs
+++ b/src/nn/sample.rs
@@ -1,14 +1,10 @@
-//! Drawing a token from a model's output logits.
-//!
-//! Every generation loop ends the same way: one position's logits become one
-//! token id. [`Sampler`] owns that step — temperature, top-k, and nucleus
-//! (top-p) filtering — and draws through [`crate::rng`], so a run under
-//! [`manual_seed`](crate::rng::manual_seed) generates reproducibly.
+//! Drawing a token from logits. Draws through [`crate::rng`], so
+//! [`manual_seed`](crate::rng::manual_seed) makes generation reproducible.
use crate::rng;
use crate::tensor::Tensor;
-/// Turns one position's logits into a token id.
+/// Turns one position's logits into a token id: temperature, top-k, top-p.
///
/// ```
/// use fastnn::nn::Sampler;
@@ -27,7 +23,7 @@ pub struct Sampler {
}
impl Sampler {
- /// Plain sampling from the softmax of the logits.
+ /// Plain softmax sampling.
pub fn new() -> Sampler {
Sampler {
temperature: 1.0,
@@ -37,12 +33,12 @@ impl Sampler {
}
}
- /// Always the most likely token. Equivalent to `temperature(0.0)`.
+ /// Always the most likely token.
pub fn greedy() -> Sampler {
Sampler::new().temperature(0.0)
}
- /// Below 1.0 sharpens the distribution, above 1.0 flattens it; 0.0 is greedy.
+ /// Below 1.0 sharpens, above flattens, 0.0 is greedy.
pub fn temperature(mut self, temperature: f32) -> Sampler {
assert!(
temperature >= 0.0,
@@ -66,8 +62,7 @@ impl Sampler {
self
}
- /// Discourage tokens that already appeared. 1.0 is off; 1.1–1.3 is the
- /// usual range. Applied through [`sample_with_history`](Self::sample_with_history).
+ /// Discourage tokens already seen. 1.0 is off. Needs [`sample_with_history`](Self::sample_with_history).
pub fn repetition_penalty(mut self, penalty: f32) -> Sampler {
assert!(
penalty >= 1.0,
@@ -77,19 +72,13 @@ impl Sampler {
self
}
- /// Draw a token id from `logits`, the unnormalized scores for one position.
- ///
- /// Any shape holding exactly the vocabulary — `[vocab]` or `[1, vocab]` —
- /// is fine; the elements are read in order.
+ /// Draw a token id from one position's logits, any shape holding the vocabulary.
pub fn sample(&self, logits: &Tensor) -> usize {
self.sample_with_history(logits, &[])
}
- /// [`sample`](Self::sample), penalising tokens that appear in `history`.
- ///
- /// A positive logit is divided by the penalty and a negative one multiplied,
- /// so "already used" always means "less likely", whichever side of zero the
- /// score sits on.
+ /// [`sample`](Self::sample), penalising tokens in `history`. Positive logits are
+ /// divided and negative ones multiplied, so both get less likely.
pub fn sample_with_history(&self, logits: &Tensor, history: &[usize]) -> usize {
let mut logits = logits.to_vec();
assert!(!logits.is_empty(), "cannot sample from empty logits");
@@ -106,8 +95,7 @@ impl Sampler {
}
}
- // Most likely first. Sorting the whole vocabulary is fine at the sizes
- // sampling sees, and gives top-k and top-p the same prefix structure.
+ // most likely first, so top-k and top-p are both prefixes
let mut order: Vec = (0..logits.len()).collect();
order.sort_by(|&a, &b| logits[b].total_cmp(&logits[a]));
@@ -118,7 +106,7 @@ impl Sampler {
order.truncate(k);
}
- // Softmax over the survivors, subtracting the max so exp cannot overflow.
+ // subtract the max so exp can't overflow
let max = logits[order[0]] / self.temperature;
let mut probs: Vec = order
.iter()
@@ -142,7 +130,7 @@ impl Sampler {
probs.iter_mut().for_each(|q| *q /= sum);
}
- // Inverse-CDF draw; the fallback covers rounding leaving a sliver above 1.
+ // the fallback covers rounding leaving the sum just under 1
let mut draw = rng::uniform(1, 0.0, 1.0)[0];
for (&index, &prob) in order.iter().zip(&probs) {
draw -= prob;
diff --git a/src/nn/sequential.rs b/src/nn/sequential.rs
index b6dbac9..bc7dc67 100644
--- a/src/nn/sequential.rs
+++ b/src/nn/sequential.rs
@@ -5,7 +5,7 @@ use crate::tensor::Tensor;
use super::module::{scoped, Module};
use super::param::{Buffer, Param};
-/// Runs its layers in order, feeding each output into the next.
+/// Runs its layers in order.
///
/// ```
/// use fastnn::prelude::*;
@@ -31,14 +31,13 @@ impl Sequential {
}
/// Append a layer, returning `self` so calls chain.
- // Named for the builder pattern, not arithmetic; `std::ops::Add` is unrelated.
#[allow(clippy::should_implement_trait)]
pub fn add(mut self, layer: impl Module + 'static) -> Sequential {
self.layers.push(Box::new(layer));
self
}
- /// Append a layer to an existing model.
+ /// Append a layer in place.
pub fn push(&mut self, layer: impl Module + 'static) {
self.layers.push(Box::new(layer));
}
diff --git a/src/nn/shape.rs b/src/nn/shape.rs
index 122e819..b46f51c 100644
--- a/src/nn/shape.rs
+++ b/src/nn/shape.rs
@@ -4,21 +4,18 @@ use crate::tensor::Tensor;
use super::module::Module;
-/// Collapse everything after the batch dimension into one.
-///
-/// The bridge from a convolution stack to a classifier head:
-/// `[N, C, H, W] → [N, C·H·W]`.
+/// Flatten everything after the batch dimension, `[N, C, H, W]` to `[N, C·H·W]`.
pub struct Flatten {
from: usize,
}
impl Flatten {
- /// Keep dimension 0 as the batch, flatten the rest.
+ /// Flatten from dimension 1.
pub fn new() -> Flatten {
Flatten { from: 1 }
}
- /// Keep dimensions before `from`, flatten from there on.
+ /// Flatten from dimension `from`.
pub fn from(from: usize) -> Flatten {
Flatten { from }
}
@@ -36,8 +33,7 @@ impl Module for Flatten {
}
}
-/// Reshape to a fixed shape. `-1` infers one dimension, which is how you keep a
-/// variable batch size.
+/// Reshape to a fixed shape. `-1` infers one dimension.
pub struct Reshape {
dims: Vec,
}
diff --git a/src/nn/transformer.rs b/src/nn/transformer.rs
index 6ca7393..b9dc019 100644
--- a/src/nn/transformer.rs
+++ b/src/nn/transformer.rs
@@ -26,12 +26,7 @@ impl Activation {
}
}
-/// One pre-norm transformer block: attention, then a feed-forward network, each
-/// wrapped in a residual connection.
-///
-/// Pre-norm — normalizing the *input* to each sub-layer rather than the output —
-/// leaves the residual path unnormalized from embedding to logits, which is what
-/// lets deep stacks train without a learning-rate warmup schedule tuned per depth.
+/// One pre-norm transformer block: attention then feed-forward, each in a residual.
pub struct TransformerBlock {
pub attention: MultiHeadAttention,
pub norm_attention: LayerNorm,
@@ -98,9 +93,7 @@ impl TransformerBlock {
}
}
- /// [`forward`](Module::forward) with a key padding mask: `[batch, len]`,
- /// 1 for real positions and 0 for padding, so a batch of unequal-length
- /// sequences attends only to its own content.
+ /// [`forward`](Module::forward) with a key padding mask, `[batch, len]`, 1 for real positions.
pub fn forward_masked(&self, input: &Tensor, key_mask: Option<&Tensor>) -> Tensor {
let normed = self.norm_attention.forward(input);
let attended =
@@ -113,11 +106,7 @@ impl TransformerBlock {
residual.add(&self.dropout.forward(&self.down.forward(&hidden)))
}
- /// [`forward`](Module::forward) with the attention keys and values cached
- /// for incremental decoding.
- ///
- /// Inference-only — dropout is skipped — and causal by nature: a cache only
- /// makes sense when later tokens cannot change earlier ones.
+ /// [`forward`](Module::forward) through a KV cache. Inference only, causal blocks only.
pub fn forward_cached(&self, input: &Tensor, cache: &mut KvCache) -> Tensor {
assert!(
self.causal,
@@ -179,7 +168,7 @@ impl TransformerStack {
})
}
- /// A causal decoder stack, as used by GPT-style language models.
+ /// A causal decoder stack, GPT style.
pub fn causal(
model_dim: usize,
heads: usize,
@@ -208,8 +197,7 @@ impl TransformerStack {
StackCache::new(self.blocks.len())
}
- /// [`forward`](Module::forward) with a key padding mask applied in every
- /// block. See [`TransformerBlock::forward_masked`].
+ /// [`forward`](Module::forward) with a key padding mask in every block.
pub fn forward_masked(&self, input: &Tensor, key_mask: Option<&Tensor>) -> Tensor {
let hidden = self
.blocks
@@ -218,8 +206,7 @@ impl TransformerStack {
self.norm.forward(&hidden)
}
- /// [`forward`](Module::forward) through per-block KV caches, for feeding a
- /// generation loop one token at a time. See [`super::cache`].
+ /// [`forward`](Module::forward) through per-block KV caches. See [`super::cache`].
pub fn forward_cached(&self, input: &Tensor, cache: &mut StackCache) -> Tensor {
assert_eq!(
cache.layers.len(),
diff --git a/src/optim/adagrad.rs b/src/optim/adagrad.rs
index d16ec88..474743d 100644
--- a/src/optim/adagrad.rs
+++ b/src/optim/adagrad.rs
@@ -1,15 +1,4 @@
-//! Adagrad and Adadelta: adaptive steps from accumulated gradient history.
-//!
-//! Adagrad divides by the square root of the *total* squared gradient ever
-//! seen — coordinates that update often slow down, rare ones stay fast, which
-//! is exactly right for sparse features and exactly wrong for long runs, where
-//! the ever-growing accumulator drives every step toward zero.
-//!
-//! Adadelta is the repair: replace the total with an exponential average, and
-//! replace the global learning rate with a running average of the updates
-//! themselves, so the step size carries the *units* of the parameter — the
-//! ratio `√E[Δx²]/√E[g²]` is dimensionally an x per g, where Adagrad's
-//! `1/√E[g²]` is not.
+//! Adagrad and Adadelta.
use crate::error::Result;
use crate::nn::Param;
@@ -19,7 +8,7 @@ use super::adam::blend;
use super::state::place;
use super::{Optimizer, OptimizerState};
-/// Adagrad: per-coordinate steps that shrink with accumulated use.
+/// Adagrad. Steps shrink as squared gradients accumulate.
pub struct Adagrad {
params: Vec,
accumulator: Vec>,
@@ -55,8 +44,6 @@ impl Optimizer for Adagrad {
}
let Some(grad) = param.grad() else { continue };
- // The accumulator is a *sum*, not an average — that monotone growth
- // is Adagrad's defining property and its known flaw.
let accumulated = match &self.accumulator[index] {
Some(prior) => prior.add(&grad.square()),
None => grad.square(),
@@ -93,10 +80,7 @@ impl Optimizer for Adagrad {
}
}
-/// Adadelta: Adagrad with a window, and no learning rate to tune.
-///
-/// `lr` survives as a final multiplier for the rare case it helps; the
-/// canonical method uses 1.0, and that is the default.
+/// Adadelta: Adagrad with an exponential window. `lr` defaults to 1.0.
pub struct Adadelta {
params: Vec,
square_avg: Vec>,
@@ -148,8 +132,7 @@ impl Optimizer for Adadelta {
let square = blend(self.square_avg[index].as_ref(), &grad.square(), self.rho);
- // The eps *inside* both roots doubles as the bootstrap: on the very
- // first step E[Δx²] is zero and the ratio would otherwise be too.
+ // eps inside the roots also bootstraps step one, where E[dx^2] is zero.
let previous_delta = self.delta_avg[index].as_ref();
let numerator = match previous_delta {
Some(delta) => delta.add_scalar(self.eps).sqrt(),
@@ -197,9 +180,6 @@ impl Optimizer for Adadelta {
mod tests {
use super::*;
- /// Two hand-computed Adagrad steps with g = 2 each time and eps = 0:
- /// step 1 divides by √4 = 2, step 2 by √8, so the second step is smaller
- /// by exactly √2.
#[test]
fn adagrad_steps_shrink_by_the_accumulated_root() {
let param = Param::new(Tensor::from_vec(vec![0.0], &[1]));
@@ -222,15 +202,13 @@ mod tests {
);
}
- /// Adadelta's step size is self-tuning: with lr = 1 and a constant
- /// gradient it settles near √eps-scale moves that keep making progress.
#[test]
fn adadelta_makes_progress_without_a_learning_rate() {
let param = Param::new(Tensor::from_vec(vec![5.0], &[1]));
let mut opt = Adadelta::new(vec![param.clone()]);
for _ in 0..200 {
- // d/dx (x²/2) = x: minimize a bowl centred at zero.
+ // gradient of x^2 / 2
param.set_grad(param.value());
opt.step();
}
diff --git a/src/optim/adam.rs b/src/optim/adam.rs
index d3dd031..837e884 100644
--- a/src/optim/adam.rs
+++ b/src/optim/adam.rs
@@ -7,7 +7,7 @@ use crate::tensor::Tensor;
use super::state::place;
use super::{Optimizer, OptimizerState};
-/// The moment estimates and update rule both optimizers share.
+/// Moment estimates shared by Adam and AdamW.
struct Moments {
first: Vec>,
second: Vec>,
@@ -29,11 +29,7 @@ impl Moments {
}
}
- /// How far to move parameter `index` given its gradient.
- ///
- /// Both moments start at zero and are therefore biased toward it early on;
- /// dividing by `1 - βᵗ` corrects for that, which is what makes the first
- /// few steps usable rather than vanishingly small.
+ /// The unscaled step for parameter `index`, bias corrected.
fn update(&mut self, index: usize, grad: &Tensor) -> Tensor {
let first = blend(self.first[index].as_ref(), grad, self.beta1);
let second = blend(self.second[index].as_ref(), &grad.square(), self.beta2);
@@ -63,9 +59,7 @@ impl Moments {
}
}
-/// `β·previous + (1-β)·value`, starting from zero on the first step.
-///
-/// The exponential moving average every adaptive optimizer is built on.
+/// Exponential moving average: `β·previous + (1-β)·value`.
pub(crate) fn blend(previous: Option<&Tensor>, value: &Tensor, beta: f32) -> Tensor {
let fresh = value.mul_scalar(1.0 - beta);
match previous {
@@ -74,7 +68,7 @@ pub(crate) fn blend(previous: Option<&Tensor>, value: &Tensor, beta: f32) -> Ten
}
}
-/// Adam: per-parameter step sizes from running gradient moments.
+/// Adam.
///
/// ```
/// # use fastnn::prelude::*;
@@ -98,21 +92,20 @@ impl Adam {
}
}
- /// Decay rates for the first and second moment. Lower `beta2` reacts faster
- /// to changing gradient scale, which transformers often want.
+ /// Decay rates for the first and second moment.
pub fn betas(mut self, beta1: f32, beta2: f32) -> Adam {
self.moments.beta1 = beta1;
self.moments.beta2 = beta2;
self
}
- /// Floor on the denominator, guarding against division by a tiny variance.
+ /// Added to the denominator for stability.
pub fn eps(mut self, eps: f32) -> Adam {
self.moments.eps = eps;
self
}
- /// L2 penalty folded into the gradient. For true weight decay use [`AdamW`].
+ /// L2 penalty added to the gradient. Use [`AdamW`] for decoupled decay.
pub fn weight_decay(mut self, decay: f32) -> Adam {
self.weight_decay = decay;
self
@@ -163,10 +156,8 @@ impl Optimizer for Adam {
/// Adam with decoupled weight decay.
///
-/// Adam's L2 penalty goes through the moment estimates, so a parameter with
-/// consistently small gradients gets decayed *harder* than one with large ones —
-/// the opposite of the intent. AdamW shrinks the weight directly instead, which
-/// is why it is the default for transformers.
+/// Adam's L2 penalty goes through the moments and decays small-gradient
+/// weights harder. AdamW shrinks the weight directly.
pub struct AdamW {
params: Vec,
moments: Moments,
@@ -210,7 +201,6 @@ impl Optimizer for AdamW {
}
let Some(grad) = param.grad() else { continue };
- // Shrink the weight first, untouched by the moment estimates.
let value = param.value().mul_scalar(1.0 - self.lr * self.weight_decay);
let update = self.moments.update(index, &grad);
param.set_value(value.sub(&update.mul_scalar(self.lr)));
diff --git a/src/optim/ema.rs b/src/optim/ema.rs
index 09765ea..4dbd20d 100644
--- a/src/optim/ema.rs
+++ b/src/optim/ema.rs
@@ -1,15 +1,9 @@
-//! An exponential moving average of the weights themselves.
-//!
-//! SGD's iterates orbit the minimum rather than landing on it; their running
-//! average sits closer to the centre of that orbit (Polyak averaging). Kept as
-//! a shadow copy with decay d, updated after each optimizer step, the average
-//! is what you *evaluate and ship* while the raw weights keep training —
-//! standard practice for diffusion models and GANs, cheap insurance elsewhere.
+//! Exponential moving average of the weights, for evaluating and shipping.
use crate::nn::Param;
use crate::tensor::Tensor;
-/// Shadow weights: `shadow ← d·shadow + (1−d)·param` after every step.
+/// Shadow weights: `shadow = d * shadow + (1 - d) * param` after every step.
///
/// ```
/// # use fastnn::prelude::*;
@@ -33,14 +27,13 @@ pub struct Ema {
}
impl Ema {
- /// Typical decays are 0.99–0.9999; higher = smoother but slower to follow.
+ /// `decay` is usually between 0.99 and 0.9999.
pub fn new(params: Vec, decay: f32) -> Ema {
assert!(
(0.0..1.0).contains(&decay),
"decay must be in [0, 1), got {decay}"
);
- // Seeding the shadow at the current weights (not zero) avoids a long
- // bias toward the origin that 1/(1−d) steps would otherwise carry.
+ // Seeded at the current weights, not zero, to avoid a long bias toward the origin.
let shadow = params.iter().map(|p| p.value().detach()).collect();
Ema {
params,
@@ -59,11 +52,7 @@ impl Ema {
}
}
- /// Exchange model weights and averaged weights.
- ///
- /// Symmetric on purpose: one call to evaluate the average, a second to put
- /// the raw training weights back. Being a swap rather than a copy, nothing
- /// is lost if the calls are unbalanced — the other set is always held here.
+ /// Swap the model's weights with the averaged ones. Call again to swap back.
pub fn swap(&mut self) {
for (shadow, param) in self.shadow.iter_mut().zip(&self.params) {
let raw = param.value().detach();
@@ -72,7 +61,7 @@ impl Ema {
}
}
- /// The averaged tensors, in parameter order — for saving a shipped model.
+ /// The averaged tensors in parameter order.
pub fn averaged(&self) -> &[Tensor] {
&self.shadow
}
@@ -82,8 +71,6 @@ impl Ema {
mod tests {
use super::*;
- /// Hand-computed: shadow seeds at 1, weight jumps to 3, one update with
- /// d = 0.9 gives 0.9·1 + 0.1·3 = 1.2.
#[test]
fn update_follows_the_definition_exactly() {
let param = Param::new(Tensor::from_vec(vec![1.0], &[1]));
diff --git a/src/optim/lion.rs b/src/optim/lion.rs
index 90c20a6..7ebf247 100644
--- a/src/optim/lion.rs
+++ b/src/optim/lion.rs
@@ -1,15 +1,7 @@
-//! Lion: sign-of-momentum updates (Chen et al., 2023).
+//! Lion (Chen et al., 2023): updates by the sign of an interpolated momentum.
//!
-//! Lion keeps one moment instead of Adam's two and updates with only the
-//! *sign* of an interpolated momentum, so every coordinate moves by exactly
-//! ±lr. That uniform magnitude acts as a strong implicit regulariser and
-//! halves optimizer memory; the price is that lr and weight decay need
-//! retuning (typically lr ~3–10× smaller than AdamW's, decay ~3–10× larger,
-//! keeping their product roughly constant).
-//!
-//! The update uses one β for the *step* direction and another for the stored
-//! momentum — that asymmetry (β₁ < β₂) is the discovered detail that makes
-//! Lion more than sign-SGD-with-momentum.
+//! Every coordinate moves by exactly lr, so lr and weight decay need retuning
+//! compared to AdamW.
use crate::error::Result;
use crate::nn::Param;
@@ -19,7 +11,9 @@ use super::adam::blend;
use super::state::place;
use super::{Optimizer, OptimizerState};
-/// Lion. Memory-light, uniformly sized steps.
+const SIGN_EPS: f32 = 1e-30;
+
+/// Lion.
pub struct Lion {
params: Vec,
momentum: Vec>,
@@ -50,18 +44,16 @@ impl Lion {
self
}
- /// Decoupled weight decay, applied at the current learning rate.
+ /// Decoupled weight decay.
pub fn weight_decay(mut self, decay: f32) -> Lion {
self.weight_decay = decay;
self
}
}
-/// `sign(x)` as `x/(|x|+ε)`: exactly 0 at 0, and within ε/|x| of ±1 elsewhere —
-/// far below f32 resolution for any gradient that matters. Staying inside the
-/// existing element-wise ops keeps the update on the parameters' device.
+/// `x/(|x|+eps)`, built from element-wise ops so it stays on the parameters' device.
fn sign(x: &Tensor) -> Tensor {
- x.div(&x.abs().add_scalar(1e-30))
+ x.div(&x.abs().add_scalar(SIGN_EPS))
}
impl Optimizer for Lion {
@@ -74,10 +66,9 @@ impl Optimizer for Lion {
let Some(grad) = param.grad() else { continue };
let value = param.value();
- // Direction: a *short-memory* interpolation toward the fresh gradient.
let previous = self.momentum[index].as_ref();
let direction = sign(&blend(previous, &grad, self.beta1));
- // Stored momentum: a *long-memory* average, updated independently.
+ // The step uses beta1 and the stored momentum uses beta2.
self.momentum[index] = Some(blend(previous, &grad, self.beta2));
let decayed = if self.weight_decay == 0.0 {
@@ -118,7 +109,6 @@ impl Optimizer for Lion {
mod tests {
use super::*;
- /// Every coordinate moves by exactly ±lr, whatever the gradient magnitude.
#[test]
fn steps_are_uniform_regardless_of_gradient_scale() {
let param = Param::new(Tensor::from_vec(vec![0.0, 0.0, 0.0], &[3]));
@@ -142,7 +132,6 @@ mod tests {
);
}
- /// A zero gradient with zero momentum moves nothing: sign(0) = 0.
#[test]
fn zero_gradient_takes_no_step() {
let param = Param::new(Tensor::from_vec(vec![1.0], &[1]));
diff --git a/src/optim/lookahead.rs b/src/optim/lookahead.rs
index 26dfd94..7e98579 100644
--- a/src/optim/lookahead.rs
+++ b/src/optim/lookahead.rs
@@ -1,10 +1,4 @@
-//! Lookahead: slow weights that follow any inner optimizer (Zhang et al., 2019).
-//!
-//! The inner optimizer explores for `k` fast steps; then the *slow* copy of
-//! the weights moves a fraction α toward wherever the fast weights ended up,
-//! and the fast weights restart from there. The averaging damps the inner
-//! optimizer's oscillations without slowing its progress — variance reduction
-//! for free, at the cost of one extra weight copy.
+//! Lookahead (Zhang et al., 2019): slow weights that follow any inner optimizer.
use crate::error::Result;
use crate::nn::Param;
@@ -29,17 +23,16 @@ pub struct Lookahead {
}
impl Lookahead {
- /// Sync every `k` inner steps, moving the slow weights by `alpha` toward
- /// the fast ones. The paper's defaults are k = 5, α = 0.5.
+ /// Sync every `k` inner steps, moving the slow weights `alpha` of the way
+ /// to the fast ones.
pub fn new(inner: O, k: usize, alpha: f32) -> Lookahead {
assert!(k >= 1, "lookahead needs k >= 1");
assert!(
(0.0..=1.0).contains(&alpha),
"alpha must be in [0, 1], got {alpha}"
);
- // The slow weights start where the parameters start (φ₀ = θ₀). Seeding
- // lazily at the first sync would seed them from weights that had
- // already taken k fast steps, making that sync a no-op.
+ // Seeding lazily at the first sync would copy weights that already took
+ // k fast steps and make that sync a no-op.
let slow = inner
.parameters()
.iter()
@@ -66,12 +59,11 @@ impl Optimizer for Lookahead {
for (index, param) in self.inner.parameters().iter().enumerate() {
let fast = param.value();
- // A parameter added after a state load may have no slow copy yet.
+ // No slow copy yet if the state was loaded without one.
let slow = match &self.slow[index] {
Some(slow) => slow.add(&fast.sub(slow).mul_scalar(self.alpha)),
None => fast.clone(),
};
- // Fast weights restart from the averaged point.
param.set_value(slow.clone());
self.slow[index] = Some(slow);
}
@@ -109,8 +101,6 @@ mod tests {
use super::*;
use crate::optim::SGD;
- /// After k fast steps of −lr each, the parameter must sit at α times the
- /// fast displacement: slow-weight interpolation, computed by hand.
#[test]
fn sync_lands_exactly_at_the_interpolated_point() {
let param = Param::new(Tensor::from_vec(vec![0.0], &[1]));
@@ -120,7 +110,7 @@ mod tests {
param.set_grad(Tensor::from_vec(vec![1.0], &[1]));
opt.step();
}
- // Fast weights walked to −0.2; slow started at 0.0; sync: 0 + 0.5·(−0.2).
+ // Fast weights walk to -0.2, so the sync lands at 0.5 * -0.2.
let got = param.value().to_vec()[0];
assert!(
(got - -0.1).abs() < 1e-6,
diff --git a/src/optim/mod.rs b/src/optim/mod.rs
index b6f04ec..0f1c5d6 100644
--- a/src/optim/mod.rs
+++ b/src/optim/mod.rs
@@ -1,7 +1,6 @@
//! Optimizers and learning-rate schedules.
//!
-//! An optimizer owns [`Param`] handles, which are shared slots rather than
-//! copies. That is what lets the training loop read:
+//! An optimizer owns shared [`Param`] handles, so `step()` needs no model borrow.
//!
//! ```
//! use fastnn::prelude::*;
@@ -14,12 +13,6 @@
//! loss.backward();
//! opt.step();
//! ```
-//!
-//! with no borrow of the model at `step()` time and no parameter list threaded
-//! through every call.
-//!
-//! Updates run on whichever device the parameters live on: they are written with
-//! tensor ops, so a GPU model never round-trips to the host to take a step.
pub mod adagrad;
pub mod adam;
@@ -46,12 +39,12 @@ pub use state::OptimizerState;
use crate::error::Result;
use crate::nn::Param;
+const NORM_EPS: f32 = 1e-6;
+
/// What every optimizer can do.
pub trait Optimizer {
- /// Apply one update using the gradients currently in each parameter.
- ///
- /// Parameters with no gradient are skipped, so a model with an unused branch
- /// still steps cleanly.
+ /// Apply one update from the gradients currently in each parameter.
+ /// Parameters with no gradient are skipped.
fn step(&mut self);
/// The parameters this optimizer updates.
@@ -60,20 +53,17 @@ pub trait Optimizer {
/// The current learning rate.
fn lr(&self) -> f32;
- /// Set the learning rate — how a [`LrSchedule`] is applied.
+ /// Set the learning rate. This is how a [`LrSchedule`] is applied.
fn set_lr(&mut self, lr: f32);
- /// Momentum, moment estimates, and step count, ready to be written to disk.
- ///
- /// Use [`save_training`](crate::serialize::save_training) rather than calling
- /// this directly unless you are storing the state somewhere of your own.
+ /// Momentum, moment estimates, and step count, for
+ /// [`save_training`](crate::serialize::save_training).
fn state(&self) -> OptimizerState;
- /// Restore what [`state`](Optimizer::state) captured.
+ /// Restore what [`state`](Optimizer::state) returned.
fn load_state(&mut self, state: OptimizerState) -> Result<()>;
- /// Clear every gradient. Call before `backward()`, or skip it to accumulate
- /// gradients across several micro-batches.
+ /// Clear every gradient. Skip it to accumulate across micro-batches.
fn zero_grad(&self) {
for param in self.parameters() {
param.zero_grad();
@@ -82,12 +72,8 @@ pub trait Optimizer {
/// Whether every parameter's gradient is finite.
///
- /// A cheap guard before `step()`: one inf in one gradient becomes an inf
- /// weight, and from there every activation downstream is NaN. Skipping the
- /// step costs one batch; taking it costs the run.
- ///
- /// Reads every gradient, so on a GPU this downloads them. Worth it around a
- /// known-unstable phase; not worth it every step of a healthy run.
+ /// Guard for `step()`: one inf gradient becomes an inf weight and then NaN
+ /// everywhere. Reads every gradient, so on a GPU it downloads them.
fn gradients_are_finite(&self) -> bool {
self.parameters()
.iter()
@@ -98,8 +84,7 @@ pub trait Optimizer {
/// Scale gradients down so their combined L2 norm is at most `max_norm`.
///
-/// Returns the norm *before* clipping, which is worth logging: a sudden spike is
-/// usually the first sign of divergence. Call between `backward()` and `step()`.
+/// Returns the norm before clipping, which is worth logging.
pub fn clip_grad_norm(params: &[Param], max_norm: f32) -> f32 {
let total: f32 = params
.iter()
@@ -109,8 +94,7 @@ pub fn clip_grad_norm(params: &[Param], max_norm: f32) -> f32 {
.sqrt();
if total > max_norm {
- // The epsilon keeps the scale finite if the norm underflows to zero.
- let scale = max_norm / (total + 1e-6);
+ let scale = max_norm / (total + NORM_EPS);
for param in params {
if let Some(grad) = param.grad() {
param.set_grad(grad.mul_scalar(scale));
@@ -122,9 +106,7 @@ pub fn clip_grad_norm(params: &[Param], max_norm: f32) -> f32 {
/// Clamp every gradient element into `[-max_value, max_value]`.
///
-/// Cruder than [`clip_grad_norm`] — it changes the gradient's direction, not
-/// just its length — but robust to a single exploded element. Call between
-/// `backward()` and `step()`.
+/// Unlike [`clip_grad_norm`] this changes the gradient's direction.
pub fn clip_grad_value(params: &[Param], max_value: f32) {
for param in params {
if let Some(grad) = param.grad() {
diff --git a/src/optim/radam.rs b/src/optim/radam.rs
index 42e9a50..872b8a5 100644
--- a/src/optim/radam.rs
+++ b/src/optim/radam.rs
@@ -1,19 +1,7 @@
-//! RAdam: Adam with its warmup requirement solved in closed form.
+//! RAdam (Liu et al., 2020): Adam with a rectified step, so it needs no warmup.
//!
-//! Adam's variance estimate `v` is built from very few samples early in
-//! training, so `1/√v̂` is itself high-variance and the first updates can be
-//! wildly mis-scaled — the practical reason Adam is run behind a warmup
-//! schedule. Liu et al. (2020) computed that variance analytically: tracking
-//! the effective sample size ρ_t of the exponential average, the adaptive step
-//! can be *rectified* by
-//!
-//! ```text
-//! r_t = √( (ρ_t−4)(ρ_t−2)ρ∞ / ((ρ∞−4)(ρ∞−2)ρ_t) ), ρ∞ = 2/(1−β₂) − 1
-//! ```
-//!
-//! which is exactly the factor that makes the adaptive step's variance match
-//! its late-training value. While ρ_t ≤ 4 the variance is undefined — too few
-//! samples — and RAdam falls back to plain momentum SGD for those steps.
+//! While the variance estimate has too few samples (rho_t <= 4) it falls back
+//! to momentum SGD.
use crate::error::Result;
use crate::nn::Param;
@@ -69,8 +57,7 @@ impl RAdam {
self
}
- /// The rectification factor for step `t`, or `None` while the variance of
- /// the adaptive step is still undefined (ρ_t ≤ 4).
+ /// The rectification factor for step `t`, or `None` while rho_t <= 4.
fn rectifier(&self, t: u64) -> Option {
let rho_inf = 2.0 / (1.0 - self.beta2) - 1.0;
let beta2_t = self.beta2.powi(t as i32);
@@ -107,7 +94,6 @@ impl Optimizer for RAdam {
mean.div(&variance.sqrt().add_scalar(self.eps))
.mul_scalar(r)
}
- // Too few samples to trust the variance: momentum SGD instead.
None => mean,
};
@@ -154,9 +140,7 @@ impl Optimizer for RAdam {
mod tests {
use super::*;
- /// With β₂ = 0.999: ρ∞ = 1999, and ρ_t ≤ 4 for the first four steps — the
- /// published fallback window. Step 5 must rectify, and with r < 1: the
- /// rectified step is *smaller* than plain Adam's would be.
+ /// With beta2 = 0.999 the first four steps fall back, step 5 rectifies.
#[test]
fn rectifier_switches_on_at_step_five_and_damps() {
let opt = RAdam::new(vec![], 1e-3);
@@ -168,7 +152,6 @@ mod tests {
}
let r = opt.rectifier(5).expect("step 5 should rectify");
assert!(r > 0.0 && r < 1.0, "early rectifier should damp, got {r}");
- // Late in training the correction must vanish: r → 1.
let late = opt.rectifier(100_000).unwrap();
assert!(
(late - 1.0).abs() < 0.01,
@@ -176,8 +159,6 @@ mod tests {
);
}
- /// The first step is pure bias-corrected momentum: with β₁ = 0.9 and g = 2,
- /// m̂₁ = g exactly, so the parameter moves by lr·g.
#[test]
fn first_step_is_exactly_momentum_sgd() {
let param = Param::new(Tensor::from_vec(vec![1.0], &[1]));
diff --git a/src/optim/rmsprop.rs b/src/optim/rmsprop.rs
index e2800c2..60af625 100644
--- a/src/optim/rmsprop.rs
+++ b/src/optim/rmsprop.rs
@@ -1,11 +1,4 @@
//! RMSprop: divide each gradient by a running estimate of its own scale.
-//!
-//! The second-moment average `E[g²]` tracks how large each coordinate's
-//! gradients have recently been, and the update `g / √E[g²]` makes every
-//! coordinate move at roughly the learning rate regardless of that scale.
-//! Adam is this plus first-moment momentum and bias correction; RMSprop's
-//! uncorrected form remains standard in reinforcement learning, where
-//! non-stationary gradients make Adam's long memory a liability.
use crate::error::Result;
use crate::nn::Param;
@@ -52,7 +45,7 @@ impl RMSprop {
}
}
- /// Decay rate of the squared-gradient average. Closer to 1 remembers longer.
+ /// Decay rate of the squared-gradient average.
pub fn alpha(mut self, alpha: f32) -> RMSprop {
assert!(
(0.0..1.0).contains(&alpha),
@@ -67,16 +60,13 @@ impl RMSprop {
self
}
- /// Classical momentum applied to the scaled update.
+ /// Momentum applied to the scaled update.
pub fn momentum(mut self, momentum: f32) -> RMSprop {
self.momentum = momentum;
self
}
- /// Subtract the squared mean from the squared average, so the denominator
- /// estimates the gradient's *variance* rather than its raw magnitude — a
- /// consistently large gradient then steps at full size instead of being
- /// damped by its own mean.
+ /// Normalise by the gradient's variance instead of its raw second moment.
pub fn centered(mut self) -> RMSprop {
self.centered = true;
self
@@ -94,8 +84,7 @@ impl Optimizer for RMSprop {
let square = blend(self.square_avg[index].as_ref(), &grad.square(), self.alpha);
- // E[g²] − E[g]² ≥ 0 mathematically, but floating point can dip a
- // hair below zero; the eps inside the root keeps sqrt defined.
+ // The variance can dip below zero in floating point, so eps goes inside the root.
let denominator = if self.centered {
let mean = blend(self.grad_avg[index].as_ref(), &grad, self.alpha);
let variance = square.sub(&mean.square());
@@ -155,8 +144,6 @@ impl Optimizer for RMSprop {
mod tests {
use super::*;
- /// One hand-computed step: g = 3, α = 0.5 ⇒ E[g²] = 4.5,
- /// update = lr·g/(√4.5 + eps) with eps = 0 for exactness.
#[test]
fn first_step_matches_the_formula_exactly() {
let param = Param::new(Tensor::from_vec(vec![1.0], &[1]));
@@ -173,10 +160,8 @@ mod tests {
);
}
- /// Centered RMSprop with a constant gradient: the variance estimate
- /// `E[g²] − E[g]²` = (1−αᵗ)αᵗ decays toward zero once the mean catches up,
- /// so late steps dwarf early ones — the signature that separates the
- /// centered form from the plain one, whose denominator settles at |g|.
+ /// With a constant gradient the variance estimate shrinks, so late steps
+ /// dwarf early ones. The plain form settles at |g| instead.
#[test]
fn centered_variant_accelerates_on_constant_gradients() {
let param = Param::new(Tensor::from_vec(vec![0.0], &[1]));
diff --git a/src/optim/schedule.rs b/src/optim/schedule.rs
index dfe245d..b3c9da2 100644
--- a/src/optim/schedule.rs
+++ b/src/optim/schedule.rs
@@ -1,8 +1,7 @@
//! Learning-rate schedules.
//!
-//! A schedule is a pure function of the step number, not something that mutates
-//! an optimizer. That makes it trivial to test, to plot, and to resume from a
-//! checkpoint — and it composes, so [`Warmup`] can wrap any of the others.
+//! A schedule is a pure function of the step, so it composes and survives a
+//! resume. [`Warmup`] wraps any other.
//!
//! ```
//! use fastnn::prelude::*;
@@ -36,14 +35,12 @@ impl LrSchedule for Constant {
/// Multiply by `gamma` every `every` steps.
pub struct StepDecay {
pub base: f32,
- /// The positive number of steps between rate changes.
+ /// Steps between rate changes, must be positive.
pub every: usize,
pub gamma: f32,
}
impl StepDecay {
- /// Create a schedule that decays every `every` steps.
- ///
/// # Panics
///
/// Panics if `every` is zero.
@@ -63,10 +60,7 @@ impl LrSchedule for StepDecay {
}
}
-/// Follow half a cosine from `base` down to `floor` over `total` steps.
-///
-/// The slow start and slow finish are the point: most of the decay happens in
-/// the middle, leaving a long low-rate tail to settle in.
+/// Half a cosine from `base` down to `floor` over `total` steps.
pub struct CosineAnnealing {
pub base: f32,
pub floor: f32,
@@ -96,10 +90,7 @@ impl LrSchedule for CosineAnnealing {
}
}
-/// Ramp linearly from zero over the first `steps`, then hand off to `inner`.
-///
-/// Adaptive optimizers have unreliable variance estimates for their first few
-/// hundred steps; warming up stops that turning into a large, badly aimed update.
+/// Ramp linearly from zero over the first `steps`, then defer to `inner`.
pub struct Warmup {
inner: S,
steps: usize,
@@ -122,24 +113,28 @@ impl LrSchedule for Warmup {
}
}
-/// One cycle: ramp up to `peak`, then cosine down well below the start.
+/// Ramp up to `peak`, then cosine down well below the start.
pub struct OneCycle {
pub peak: f32,
pub total: usize,
- /// The finite fraction of steps used for warmup, in the inclusive range [0, 1].
+ /// Fraction of steps spent warming up, in [0, 1].
pub warmup_fraction: f32,
pub start_divisor: f32,
pub final_divisor: f32,
}
+const DEFAULT_WARMUP_FRACTION: f32 = 0.3;
+const DEFAULT_START_DIVISOR: f32 = 25.0;
+const DEFAULT_FINAL_DIVISOR: f32 = 10_000.0;
+
impl OneCycle {
pub fn new(peak: f32, total: usize) -> OneCycle {
OneCycle {
peak,
total,
- warmup_fraction: 0.3,
- start_divisor: 25.0,
- final_divisor: 10_000.0,
+ warmup_fraction: DEFAULT_WARMUP_FRACTION,
+ start_divisor: DEFAULT_START_DIVISOR,
+ final_divisor: DEFAULT_FINAL_DIVISOR,
}
}
}
@@ -154,7 +149,7 @@ impl LrSchedule for OneCycle {
"warmup_fraction must be in [0, 1], got {}",
self.warmup_fraction
);
- // Converting total to f32 can round warmup above total even at a fraction of 1.
+ // total as f32 can round up, pushing warmup past total at a fraction of 1.
let warmup = ((self.warmup_fraction * self.total as f32) as usize).min(self.total);
let start = self.peak / self.start_divisor;
let end = self.peak / self.final_divisor;
@@ -210,7 +205,7 @@ mod tests {
StepDecay::new(1.0, 0, 0.5);
}
- // The fields are public, so the check in lr_at has to catch this too.
+ // The field is public, so lr_at has to check too.
#[test]
#[should_panic(expected = "every")]
fn step_decay_rejects_zero_every_set_after_new() {
@@ -272,8 +267,7 @@ mod tests {
close(schedule.lr_at(2), 0.0001);
}
- // 2^24 + 3 rounds up in f32, so a fraction of exactly 1 used to put the
- // warmup past total and underflow the subtraction after it.
+ // 2^24 + 3 rounds up in f32, which used to underflow total - warmup.
#[test]
fn one_cycle_survives_a_total_that_rounds_up_in_f32() {
let total = 16_777_219;
diff --git a/src/optim/sgd.rs b/src/optim/sgd.rs
index 5b62f3a..a93f1db 100644
--- a/src/optim/sgd.rs
+++ b/src/optim/sgd.rs
@@ -7,7 +7,7 @@ use crate::tensor::Tensor;
use super::state::place;
use super::{Optimizer, OptimizerState};
-/// Plain SGD, optionally with momentum, Nesterov lookahead, and weight decay.
+/// SGD with optional momentum, Nesterov, and weight decay.
///
/// ```
/// # use fastnn::prelude::*;
@@ -41,26 +41,25 @@ impl SGD {
}
}
- /// Carry a fraction of the previous update forward. 0.9 is the usual choice.
+ /// Fraction of the previous update carried forward.
pub fn momentum(mut self, momentum: f32) -> SGD {
self.momentum = momentum;
self
}
- /// Damp how much of each new gradient enters the velocity.
+ /// Damping on each new gradient entering the velocity.
pub fn dampening(mut self, dampening: f32) -> SGD {
self.dampening = dampening;
self
}
- /// L2 penalty, folded into the gradient.
+ /// L2 penalty added to the gradient.
pub fn weight_decay(mut self, decay: f32) -> SGD {
self.weight_decay = decay;
self
}
- /// Apply the gradient at the point momentum is heading toward rather than
- /// where it is now — a half-step of foresight that damps overshoot.
+ /// Use Nesterov momentum.
pub fn nesterov(mut self, nesterov: bool) -> SGD {
self.nesterov = nesterov;
self
@@ -77,10 +76,11 @@ impl Optimizer for SGD {
let Some(grad) = param.grad() else { continue };
let value = param.value();
- let mut grad = grad;
- if self.weight_decay != 0.0 {
- grad = grad.add(&value.mul_scalar(self.weight_decay));
- }
+ let grad = if self.weight_decay == 0.0 {
+ grad
+ } else {
+ grad.add(&value.mul_scalar(self.weight_decay))
+ };
let update = if self.momentum == 0.0 {
grad
diff --git a/src/optim/state.rs b/src/optim/state.rs
index 4d6d8f9..eb8182f 100644
--- a/src/optim/state.rs
+++ b/src/optim/state.rs
@@ -1,11 +1,9 @@
//! An optimizer's internal state, in a form that can be written to disk.
//!
-//! Saving only the weights is not enough to resume a run. Momentum and Adam's
-//! two moment estimates are what turn a raw gradient into an update, and Adam's
-//! step count drives the bias correction that divides them. Restore the weights
-//! alone and the optimizer restarts at step 1: the correction factor is at its
-//! most aggressive, so the first update after resuming is far larger than it
-//! should be and can undo a great deal of training.
+//! Optimizer state in a form that can be written to disk.
+//!
+//! Weights alone are not a resumable run: Adam's step count drives its bias
+//! correction, and restarting at step 1 makes the first update far too big.
use std::collections::BTreeMap;
@@ -13,7 +11,7 @@ use crate::error::{Error, Result};
use crate::nn::Param;
use crate::tensor::Tensor;
-/// Everything an optimizer needs to carry on exactly where it stopped.
+/// What an optimizer needs to resume where it stopped.
#[derive(Default)]
pub struct OptimizerState {
/// Per-parameter tensors, named `"/"`.
@@ -30,10 +28,7 @@ impl OptimizerState {
}
}
- /// Record one per-parameter series, such as momentum or a moment estimate.
- ///
- /// Entries still `None` — parameters that have never received a gradient —
- /// are left out, and come back as `None`.
+ /// Record one per-parameter series. `None` entries are skipped and come back as `None`.
pub fn put(&mut self, slot: &str, series: &[Option]) {
for (index, value) in series.iter().enumerate() {
if let Some(tensor) = value {
@@ -43,10 +38,8 @@ impl OptimizerState {
}
}
- /// Read back a series recorded by [`put`](OptimizerState::put).
- ///
- /// `len` is the current parameter count; a mismatch with the saved run is
- /// reported rather than silently misaligning moments onto the wrong weights.
+ /// Read back a series from [`put`](OptimizerState::put). Errors if the saved
+ /// run had more parameters than `len`.
pub fn take(&self, slot: &str, len: usize) -> Result>> {
let prefix = format!("{slot}/");
if let Some(name) = self.tensors.keys().find(|k| {
@@ -64,12 +57,8 @@ impl OptimizerState {
}
}
-/// Move a restored series onto the devices of the parameters it belongs to.
-///
-/// State is stored on the host so a checkpoint is portable between machines. The
-/// optimizer then combines it with live gradients, which sit wherever the model
-/// does — so without this a run resumed onto a GPU would hit a device mismatch on
-/// its first step.
+/// Move restored state onto each parameter's device. It is saved on the host,
+/// so skipping this breaks the first step of a run resumed on a GPU.
pub fn place(series: Vec>, params: &[Param]) -> Vec> {
series
.into_iter()
diff --git a/src/rng.rs b/src/rng.rs
index 844fa1c..7d28b86 100644
--- a/src/rng.rs
+++ b/src/rng.rs
@@ -1,7 +1,4 @@
-//! The random source behind weight init, dropout, and shuffling.
-//!
-//! One generator per thread. [`manual_seed`] reseeds the calling thread, so a
-//! single-threaded training script is reproducible end to end.
+//! The per-thread random source behind init, dropout, and shuffling.
use std::cell::RefCell;
@@ -13,12 +10,12 @@ thread_local! {
static RNG: RefCell = RefCell::new(rand::make_rng());
}
-/// Reseed this thread's generator so a run reproduces exactly.
+/// Reseed this thread's generator.
pub fn manual_seed(seed: u64) {
RNG.with(|r| *r.borrow_mut() = StdRng::seed_from_u64(seed));
}
-/// Borrow the generator. Use this for distributions the helpers below don't cover.
+/// Borrow the generator for distributions the helpers don't cover.
pub fn with_rng(f: impl FnOnce(&mut StdRng) -> T) -> T {
RNG.with(|r| f(&mut r.borrow_mut()))
}
@@ -33,12 +30,12 @@ pub fn normal(n: usize) -> Vec {
with_rng(|rng| (0..n).map(|_| rng.sample(StandardNormal)).collect())
}
-/// `n` draws from Bernoulli(`p`), as a bitmask of `true` = kept.
+/// `n` Bernoulli(`p`) draws, `true` meaning kept.
pub fn bernoulli(n: usize, p: f32) -> Vec {
with_rng(|rng| (0..n).map(|_| rng.random::() < p).collect())
}
-/// Shuffle in place, using this thread's generator so `manual_seed` covers it.
+/// Shuffle in place with this thread's generator.
pub fn shuffle(items: &mut [T]) {
use rand::seq::SliceRandom;
with_rng(|rng| items.shuffle(rng));
diff --git a/src/serialize/checkpoint.rs b/src/serialize/checkpoint.rs
index 17fb7bb..5f76b6b 100644
--- a/src/serialize/checkpoint.rs
+++ b/src/serialize/checkpoint.rs
@@ -1,6 +1,4 @@
-//! Saving and loading model weights.
-//!
-//! ## File format (`.fdl`)
+//! Saving and loading model weights in the `.fdl` format.
//!
//! ```text
//! magic "FDL\0" u32
@@ -10,13 +8,9 @@
//! name_len u32
//! name name_len bytes, UTF-8
//! rank u32
-//! shape rank × u32
-//! data numel × f32, little-endian
+//! shape rank x u32
+//! data numel x f32, little-endian
//! ```
-//!
-//! Plain and self-describing: a checkpoint carries names and shapes, so loading
-//! into a changed model reports which tensor disagrees instead of silently
-//! filling weights with the wrong numbers.
use std::collections::BTreeMap;
use std::fs::File;
@@ -30,10 +24,7 @@ use crate::tensor::Tensor;
const MAGIC: u32 = 0x4644_4C00;
const VERSION: u32 = 1;
-/// A named collection of tensors — what a saved model is.
-///
-/// `BTreeMap` so the on-disk order is stable and two checkpoints of the same
-/// model compare byte for byte.
+/// Named tensors. A `BTreeMap` keeps the file order stable.
pub type Checkpoint = BTreeMap;
/// Write a module's parameters and buffers to `path`.
@@ -57,8 +48,8 @@ pub fn save(module: &dyn Module, path: impl AsRef) -> Result<()> {
/// Read `path` into a module's parameters and buffers.
///
-/// Every parameter the module declares must be present with a matching shape;
-/// extra tensors in the file are ignored.
+/// Every parameter must be present with a matching shape. Extra tensors in the
+/// file are ignored.
pub fn load(module: &dyn Module, path: impl AsRef) -> Result<()> {
let saved = load_tensors(path)?;
@@ -140,7 +131,7 @@ pub fn load_tensors(path: impl AsRef) -> Result {
Ok(tensors)
}
-/// Look up `name` and check it is the shape the model expects.
+/// Look up `name` and check its shape.
fn take(saved: &Checkpoint, name: &str, expected: &[usize]) -> Result {
let tensor = saved
.get(name)
diff --git a/src/serialize/half.rs b/src/serialize/half.rs
index b7e2911..9487536 100644
--- a/src/serialize/half.rs
+++ b/src/serialize/half.rs
@@ -1,56 +1,39 @@
-//! Bit-exact conversion between `f32` and the two 16-bit float formats.
-//!
-//! An IEEE 754 binary16 (`f16`) keeps 1 sign, 5 exponent, and 10 mantissa
-//! bits; bfloat16 keeps 1/8/7 — the same exponent range as `f32` with less
-//! precision, which is why it never overflows where `f32` did not. Both
-//! conversions round to nearest, ties to even: the only rounding that is
-//! unbiased under accumulation, and what hardware does.
-//!
-//! Widening (`f16 → f32`) is exact — every 16-bit float is representable in
-//! `f32` — so a checkpoint stored in half precision loses information exactly
-//! once, at save time, never again on load.
-
-/// `f32` → binary16 bits, round-to-nearest-even.
+//! Conversions between `f32` and the 16-bit float formats, both round to nearest even.
+
+/// `f32` to binary16 bits.
pub fn f16_from_f32(value: f32) -> u16 {
let bits = value.to_bits();
let sign = ((bits >> 16) & 0x8000) as u16;
let exponent = ((bits >> 23) & 0xFF) as i32;
let mantissa = bits & 0x007F_FFFF;
- // Infinity and NaN: exponent all ones. A NaN must stay a NaN, so force a
- // mantissa bit on in case the payload's surviving bits are all zero.
+ // Force a mantissa bit so a NaN whose payload bits all shift out stays a NaN.
if exponent == 0xFF {
let payload = (mantissa >> 13) as u16;
return sign | 0x7C00 | if mantissa != 0 { payload | 0x0200 } else { 0 };
}
- // Rebase the exponent from f32's bias (127) to f16's (15).
let unbiased = exponent - 127;
if unbiased > 15 {
- return sign | 0x7C00; // overflows f16's range → ±∞
+ return sign | 0x7C00;
}
if unbiased < -14 {
- // Subnormal in f16: the implicit leading 1 becomes explicit and the
- // 24-bit significand shifts right until the exponent reaches −14.
+ // Subnormal in f16, so the implicit leading 1 becomes explicit.
if unbiased < -25 {
- return sign; // below half the smallest subnormal → ±0
+ return sign;
}
let significand = mantissa | 0x0080_0000;
let shift = (-14 - unbiased + 13) as u32;
let halfway = 1u32 << (shift - 1);
let truncated = significand >> shift;
- // Ties to even: round up on > halfway, or on exactly halfway when the
- // kept part is odd.
let remainder = significand & ((1 << shift) - 1);
let rounded = truncated
+ u32::from(remainder > halfway || (remainder == halfway && truncated & 1 == 1));
- // A rounded subnormal can carry into the smallest normal; the bit
- // layout makes that carry land in the exponent field on its own.
+ // A carry into the smallest normal lands in the exponent field on its own.
return sign | rounded as u16;
}
- // Normal: drop 13 mantissa bits with the same tie-to-even rule.
let halfway = 0x0000_1000u32;
let remainder = mantissa & 0x1FFF;
let truncated = mantissa >> 13;
@@ -58,7 +41,6 @@ pub fn f16_from_f32(value: f32) -> u16 {
truncated + u32::from(remainder > halfway || (remainder == halfway && truncated & 1 == 1));
let mut exponent16 = (unbiased + 15) as u32;
if rounded == 0x400 {
- // Mantissa overflowed 10 bits: 1.111…1 rounded up to 10.0…0.
rounded = 0;
exponent16 += 1;
if exponent16 >= 0x1F {
@@ -68,7 +50,7 @@ pub fn f16_from_f32(value: f32) -> u16 {
sign | ((exponent16 as u16) << 10) | rounded as u16
}
-/// binary16 bits → `f32`, exactly.
+/// binary16 bits to `f32`, exactly.
pub fn f32_from_f16(bits: u16) -> f32 {
let sign = ((bits & 0x8000) as u32) << 16;
let exponent = (bits >> 10) & 0x1F;
@@ -76,9 +58,8 @@ pub fn f32_from_f16(bits: u16) -> f32 {
match (exponent, mantissa) {
(0, 0) => f32::from_bits(sign),
- // Subnormal: value is mantissa · 2⁻²⁴, always exact in f32.
(0, m) => {
- let magnitude = m as f32 * f32::from_bits(0x3380_0000); // 2⁻²⁴
+ let magnitude = m as f32 * f32::from_bits(0x3380_0000); // 2^-24
if sign != 0 {
-magnitude
} else {
@@ -87,27 +68,22 @@ pub fn f32_from_f16(bits: u16) -> f32 {
}
(0x1F, 0) => f32::from_bits(sign | 0x7F80_0000),
(0x1F, _) => f32::NAN,
- // Normal: rebias 15 → 127 and left-align the mantissa.
_ => f32::from_bits(sign | ((exponent as u32 + 112) << 23) | (mantissa << 13)),
}
}
-/// `f32` → bfloat16 bits, round-to-nearest-even.
-///
-/// bfloat16 is the top half of an `f32`, so rounding is one addition on the
-/// raw bits: add half of the dropped range, plus one more when the kept part
-/// is odd — the carry implements ties-to-even for free.
+/// `f32` to bfloat16 bits. The addition rounds ties to even through the carry.
pub fn bf16_from_f32(value: f32) -> u16 {
let bits = value.to_bits();
if value.is_nan() {
- // Truncation could zero the payload and turn NaN into ∞; pin one bit.
+ // Truncation could turn the NaN into infinity, so pin a payload bit.
return ((bits >> 16) as u16) | 0x0040;
}
let round = 0x7FFF + ((bits >> 16) & 1);
((bits + round) >> 16) as u16
}
-/// bfloat16 bits → `f32`, exactly: the bits *are* the top half.
+/// bfloat16 bits to `f32`, exactly.
pub fn f32_from_bf16(bits: u16) -> f32 {
f32::from_bits((bits as u32) << 16)
}
@@ -135,8 +111,6 @@ mod tests {
"f16 round trip of {value}"
);
}
- // bf16 keeps only the top 16 bits, so representable means "low half
- // zero" — constructed from bits to make that property visible.
for bits in [
0x0000_0000u32,
0xBF80_0000,
@@ -152,19 +126,15 @@ mod tests {
#[test]
fn ties_round_to_even() {
- // 1 + 2⁻¹¹ sits exactly between 1.0 and 1 + 2⁻¹⁰ (f16's next value):
- // the even mantissa wins, so it rounds *down* to 1.0 —
+ // Exact ties: the even mantissa wins.
let tie_low = 1.0 + f32::from_bits(0x3A00_0000); // 2⁻¹¹
assert_eq!(f32_from_f16(f16_from_f32(tie_low)), 1.0);
- // — while 1 + 3·2⁻¹¹ sits between 1 + 2⁻¹⁰ and 1 + 2⁻⁹ and rounds *up*
- // to the even 1 + 2⁻⁹.
let tie_high = 1.0 + 3.0 * f32::from_bits(0x3A00_0000);
assert_eq!(
f32_from_f16(f16_from_f32(tie_high)),
1.0 + f32::from_bits(0x3B00_0000)
);
- // Same rule for bf16 at its own precision: 1 + 2⁻⁸ is a tie, 1.0 is even.
let tie = 1.0 + f32::from_bits(0x3B80_0000); // 2⁻⁸
assert_eq!(f32_from_bf16(bf16_from_f32(tie)), 1.0);
}
@@ -179,10 +149,8 @@ mod tests {
assert!(f32_from_f16(f16_from_f32(f32::NAN)).is_nan());
assert!(f32_from_bf16(bf16_from_f32(f32::NAN)).is_nan());
- // Values past f16's range overflow to infinity rather than wrapping.
assert_eq!(f32_from_f16(f16_from_f32(1.0e6)), f32::INFINITY);
assert_eq!(f32_from_f16(f16_from_f32(-1.0e6)), f32::NEG_INFINITY);
- // Values below half the smallest subnormal flush to signed zero.
assert_eq!(
f32_from_f16(f16_from_f32(1.0e-9)).to_bits(),
0.0f32.to_bits()
@@ -193,8 +161,6 @@ mod tests {
);
}
- /// Rounding must never be off by more than half an ulp: sweep a range of
- /// bit patterns and compare against the nearest of the two neighbours.
#[test]
fn rounding_error_is_at_most_half_an_ulp() {
for i in 0..2000u32 {
diff --git a/src/serialize/mod.rs b/src/serialize/mod.rs
index dc318f1..2c0bffe 100644
--- a/src/serialize/mod.rs
+++ b/src/serialize/mod.rs
@@ -1,13 +1,5 @@
-//! Checkpointing to disk.
-//!
-//! - [`checkpoint`] — weights, for a finished model
-//! - [`training`] — weights plus optimizer state, for resuming a run
-//! - [`safetensors`] — the interchange format the wider ecosystem shares,
-//! including reading `F16`/`BF16` checkpoints widened to `f32`
-//! - [`half`] — the bit-exact 16-bit float conversions behind that
-//!
-//! The first two use the same `.fdl` container; a training checkpoint just
-//! holds more.
+//! Checkpoints: `.fdl` weights, `.fdl` training state, safetensors, and the
+//! 16-bit float conversions safetensors needs.
pub mod checkpoint;
pub mod half;
diff --git a/src/serialize/safetensors.rs b/src/serialize/safetensors.rs
index 9d642ab..b668095 100644
--- a/src/serialize/safetensors.rs
+++ b/src/serialize/safetensors.rs
@@ -1,20 +1,11 @@
-//! The safetensors interchange format — how weights arrive from elsewhere.
+//! Reading and writing safetensors.
//!
-//! The format is deliberately simple: an 8-byte little-endian header length, a
-//! JSON header mapping each tensor name to `{dtype, shape, data_offsets}`, and
-//! the raw tensor bytes. No code execution, no pickle — which is exactly why
-//! the wider ecosystem standardized on it for distributing weights.
+//! Writes `F32`. Reads `F32`, `F16`, and `BF16`, widened to `f32`. Loading is
+//! partial: the [`LoadReport`] lists what was missing or unused.
//!
//! ```text
-//! u64 header_len │ {"w": {"dtype":"F32","shape":[2,3],"data_offsets":[0,24]}, …} │ data…
+//! u64 header_len | {"w": {"dtype":"F32","shape":[2,3],"data_offsets":[0,24]}, ...} | data
//! ```
-//!
-//! Writing always stores `F32`, FastNN's compute type. Reading additionally
-//! accepts `F16` and `BF16`, widening on load — exact, since every 16-bit
-//! float is representable in `f32`. Loading is *partial by design*: what
-//! matches is loaded, and the [`LoadReport`] names what was missing and what
-//! went unused, because a checkpoint from another framework rarely aligns
-//! perfectly on the first try and "which keys?" is the whole debugging session.
use std::collections::{BTreeMap, HashMap};
use std::path::Path;
@@ -25,13 +16,12 @@ use crate::tensor::Tensor;
use super::half::{f32_from_bf16, f32_from_f16};
-/// What a partial load actually did. Empty `missing` and `unexpected` mean the
-/// file and the model agreed exactly.
+/// What a partial load did. Empty `missing` and `unexpected` mean an exact match.
#[derive(Debug, Default)]
pub struct LoadReport {
/// Parameters and buffers that received values.
pub loaded: Vec,
- /// Model tensors the file had nothing for — they keep their current values.
+ /// Model tensors absent from the file. They keep their values.
pub missing: Vec,
/// File tensors no model tensor claimed.
pub unexpected: Vec,
@@ -39,7 +29,6 @@ pub struct LoadReport {
/// Write a module's parameters and buffers as `F32` safetensors.
pub fn save_safetensors(module: &dyn Module, path: impl AsRef) -> Result<()> {
- // BTreeMap for a stable on-disk order, as the .fdl writer does.
let mut tensors: BTreeMap = BTreeMap::new();
for (name, param) in module.named_parameters() {
tensors.insert(name, param.value().cpu());
@@ -51,7 +40,7 @@ pub fn save_safetensors(module: &dyn Module, path: impl AsRef) -> Result<(
let mut header = String::from("{");
let mut offset = 0usize;
for (index, (name, tensor)) in tensors.iter().enumerate() {
- let bytes = tensor.numel() * 4;
+ let bytes = tensor.numel() * std::mem::size_of::();
if index > 0 {
header.push(',');
}
@@ -81,17 +70,13 @@ pub fn save_safetensors(module: &dyn Module, path: impl AsRef) -> Result<(
/// Load matching tensors from a safetensors file into a module.
///
-/// Shapes must agree wherever a name matches — that is an error, not a report
-/// entry, because silently skipping a shape clash loads a model that *looks*
-/// initialized. Names that simply do not match are reported and skipped.
+/// A shape mismatch on a matching name is an error. Unmatched names are only
+/// reported.
pub fn load_safetensors(module: &dyn Module, path: impl AsRef) -> Result {
load_safetensors_renamed(module, path, &HashMap::new())
}
-/// [`load_safetensors`] through a name map: `file key → FastNN parameter path`.
-///
-/// The declarative escape hatch for foreign checkpoints — import logic becomes
-/// a table (`"transformer.h.0.attn.c_attn.weight"` → `"blocks.0.…"`), not code.
+/// [`load_safetensors`] with a map from file key to FastNN parameter path.
pub fn load_safetensors_renamed(
module: &dyn Module,
path: impl AsRef,
@@ -100,7 +85,7 @@ pub fn load_safetensors_renamed(
let bytes = std::fs::read(path)?;
let (entries, data) = parse_file(&bytes)?;
- // Index by the *effective* name: the mapping applied to the file's keys.
+ // Keyed by the name after renaming.
let mut by_name: HashMap<&str, &Entry> = HashMap::new();
for entry in &entries {
let effective = rename
@@ -154,8 +139,6 @@ pub fn load_safetensors_renamed(
Ok(report)
}
-// ── The header ───────────────────────────────────────────────────────────────
-
#[derive(Debug, PartialEq)]
enum Dtype {
F32,
@@ -205,7 +188,7 @@ fn bad(message: impl Into) -> Error {
Error::Checkpoint(format!("safetensors: {}", message.into()))
}
-/// Split a file into validated header entries and the data section.
+/// Split a file into validated entries and the data section.
fn parse_file(bytes: &[u8]) -> Result<(Vec, &[u8])> {
if bytes.len() < 8 {
return Err(bad("file shorter than its header length field"));
@@ -244,11 +227,8 @@ fn parse_file(bytes: &[u8]) -> Result<(Vec, &[u8])> {
Ok((entries, data))
}
-/// A hand-rolled parser for exactly the JSON the format specifies: one flat
-/// object of tensor entries plus an optional string-to-string `__metadata__`.
-///
-/// Small enough to read in one sitting on purpose — the alternative is a JSON
-/// dependency for one fixed schema.
+/// Minimal parser for the safetensors header: a flat object of tensor entries
+/// plus an optional string-to-string `__metadata__`.
struct Header<'a> {
bytes: &'a [u8],
pos: usize,
@@ -277,20 +257,13 @@ impl<'a> Header<'a> {
} else {
entries.push(self.entry(name)?);
}
- match self.next()? {
- b',' => continue,
- b'}' => return Ok(entries),
- other => {
- return Err(bad(format!(
- "expected ',' or '}}', found '{}'",
- other as char
- )))
- }
+ if self.list_end(b'}')? {
+ return Ok(entries);
}
}
}
- /// `{"dtype": "...", "shape": [...], "data_offsets": [a, b]}`, keys in any order.
+ /// `{"dtype": "...", "shape": [...], "data_offsets": [a, b]}` in any key order.
fn entry(&mut self, name: String) -> Result {
let (mut dtype, mut shape, mut offsets) = (None, None, None);
self.expect(b'{')?;
@@ -320,15 +293,8 @@ impl<'a> Header<'a> {
}
other => return Err(bad(format!("unknown key '{other}' in '{name}'"))),
}
- match self.next()? {
- b',' => continue,
- b'}' => break,
- other => {
- return Err(bad(format!(
- "expected ',' or '}}', found '{}'",
- other as char
- )))
- }
+ if self.list_end(b'}')? {
+ break;
}
}
match (dtype, shape, offsets) {
@@ -344,7 +310,7 @@ impl<'a> Header<'a> {
}
}
- /// `[1, 2, 3]` (or `[]`), as machine sizes.
+ /// `[1, 2, 3]` or `[]`.
fn numbers(&mut self) -> Result> {
let mut values = Vec::new();
self.expect(b'[')?;
@@ -354,15 +320,8 @@ impl<'a> Header<'a> {
}
loop {
values.push(self.number()?);
- match self.next()? {
- b',' => continue,
- b']' => return Ok(values),
- other => {
- return Err(bad(format!(
- "expected ',' or ']', found '{}'",
- other as char
- )))
- }
+ if self.list_end(b']')? {
+ return Ok(values);
}
}
}
@@ -382,7 +341,7 @@ impl<'a> Header<'a> {
.map_err(|_| bad("number too large"))
}
- /// A JSON string with the standard escapes, including `\uXXXX` pairs.
+ /// A JSON string, including `\uXXXX` surrogate pairs.
fn string(&mut self) -> Result {
self.expect(b'"')?;
let mut out = String::new();
@@ -403,7 +362,6 @@ impl<'a> Header<'a> {
},
byte if byte < 0x20 => return Err(bad("raw control character in string")),
byte => {
- // Re-assemble multi-byte UTF-8 (the header is validated UTF-8).
let width = match byte {
0x00..=0x7F => 1,
0xC0..=0xDF => 2,
@@ -420,7 +378,6 @@ impl<'a> Header<'a> {
fn unicode_escape(&mut self) -> Result {
let unit = self.hex4()?;
- // Surrogate pair: a high half must be followed by an escaped low half.
if (0xD800..0xDC00).contains(&unit) {
if self.raw_byte()? != b'\\' || self.raw_byte()? != b'u' {
return Err(bad("lone high surrogate"));
@@ -446,7 +403,7 @@ impl<'a> Header<'a> {
Ok(value)
}
- /// `{"key": "value", ...}` — the metadata object, values discarded.
+ /// Skip the `__metadata__` object.
fn skip_string_object(&mut self) -> Result<()> {
self.expect(b'{')?;
if self.peek()? == b'}' {
@@ -457,19 +414,24 @@ impl<'a> Header<'a> {
self.string()?;
self.expect(b':')?;
self.string()?;
- match self.next()? {
- b',' => continue,
- b'}' => return Ok(()),
- other => {
- return Err(bad(format!(
- "expected ',' or '}}', found '{}'",
- other as char
- )))
- }
+ if self.list_end(b'}')? {
+ return Ok(());
}
}
}
+ /// Consume a `,` (more to come, false) or `close` (done, true).
+ fn list_end(&mut self, close: u8) -> Result {
+ match self.next()? {
+ b',' => Ok(false),
+ c if c == close => Ok(true),
+ other => Err(bad(format!(
+ "expected ',' or '{}', found '{}'",
+ close as char, other as char
+ ))),
+ }
+ }
+
fn skip_whitespace(&mut self) {
while self.pos < self.bytes.len() && self.bytes[self.pos].is_ascii_whitespace() {
self.pos += 1;
@@ -490,7 +452,7 @@ impl<'a> Header<'a> {
Ok(byte)
}
- /// The next byte with no whitespace skipping — inside strings, spaces count.
+ /// The next byte without skipping whitespace, for use inside strings.
fn raw_byte(&mut self) -> Result {
let byte = self
.bytes
diff --git a/src/serialize/training.rs b/src/serialize/training.rs
index 2b1e1a4..4230135 100644
--- a/src/serialize/training.rs
+++ b/src/serialize/training.rs
@@ -1,10 +1,4 @@
-//! Checkpointing a run, not just a model.
-//!
-//! [`save`](super::save) writes weights, which is what you want for a finished
-//! model. Resuming an interrupted run needs more: the optimizer's momentum and
-//! moment estimates, its step count, and where the schedule had got to. Without
-//! them the optimizer restarts cold and the first update after resuming lands
-//! far harder than it should.
+//! Checkpoints that can resume a run: weights, optimizer state, and the step.
//!
//! ```no_run
//! # use fastnn::prelude::*;
@@ -13,7 +7,7 @@
//! # let total_steps = 10_000;
//! let mut opt = Adam::new(model.parameters(), 1e-3);
//!
-//! // Pick up where the last run stopped, or start fresh.
+//! // Resume if a checkpoint exists.
//! let mut step = load_training(&model, &mut opt, "run.fdl").unwrap_or(0);
//!
//! while step < total_steps {
@@ -35,8 +29,7 @@ use crate::tensor::Tensor;
use super::checkpoint::{load_tensors, save_tensors, Checkpoint};
-/// Namespaces inside the file. `/` cannot collide with the `.` that separates
-/// module paths in a parameter name.
+/// Name prefixes inside the file. `/` can't clash with the `.` in module paths.
const PARAM: &str = "param/";
const BUFFER: &str = "buffer/";
const OPTIM: &str = "optim/";
@@ -69,11 +62,10 @@ pub fn save_training(
save_tensors(&file, path)
}
-/// Restore what [`save_training`] wrote, returning the step it stopped at.
+/// Restore what [`save_training`] wrote and return the step it stopped at.
///
-/// Errors if `path` is a plain model checkpoint rather than a training one —
-/// resuming from weights alone would silently reset the optimizer, which is the
-/// exact failure this module exists to prevent.
+/// A plain model checkpoint is an error, since resuming from it would silently
+/// reset the optimizer.
pub fn load_training(
model: &dyn Module,
optimizer: &mut dyn Optimizer,
@@ -82,7 +74,7 @@ pub fn load_training(
let file = load_tensors(path)?;
if !file.contains_key(STEP) {
return Err(Error::Checkpoint(
- "not a training checkpoint — it has no optimizer state, so resuming \
+ "not a training checkpoint; it has no optimizer state, so resuming \
from it would restart the optimizer. Use load() if that is what you want."
.into(),
));
diff --git a/src/tensor/checked.rs b/src/tensor/checked.rs
index aa5c7e3..9a6f2b2 100644
--- a/src/tensor/checked.rs
+++ b/src/tensor/checked.rs
@@ -1,23 +1,5 @@
-//! Fallible variants of the ops that can disagree about shape.
-//!
-//! The ordinary ops panic, because inside a model a shape mismatch is a bug and
-//! threading `Result` through every `add` would bury the maths in `?`. But at the
-//! edge of an application — a shape from a config file, a batch from a user
-//! upload, a checkpoint from elsewhere — a mismatch is *data*, and an app should
-//! report it rather than die.
-//!
-//! Every function here validates, then delegates to the panicking op. There is
-//! one implementation of each operation, and this file only decides whether
-//! reaching it is safe.
-//!
-//! ```
-//! # use fastnn::prelude::*;
-//! let x = Tensor::zeros(&[4, 8]);
-//! let w = Tensor::zeros(&[3, 3]);
-//!
-//! assert!(x.try_matmul(&w).is_err()); // reported
-//! // x.matmul(&w) // would panic
-//! ```
+//! `try_*` variants of the ops that can disagree about shape. Each validates and
+//! then calls the panicking op, so there is only one implementation of each.
use crate::error::{Error, Result};
@@ -128,8 +110,6 @@ impl Tensor {
Ok(self.index_select(ids))
}
- // ── Validation ───────────────────────────────────────────────────────────
-
fn check_same_device(&self, other: &Tensor, op: &str) -> Result<()> {
if self.device() != other.device() {
return Err(Error::Shape(format!(
@@ -141,7 +121,6 @@ impl Tensor {
Ok(())
}
- /// Matmul needs rank 2+ on both sides before any inner axis even exists.
fn check_matmul_rank(&self, other: &Tensor) -> Result<()> {
if self.ndim() < 2 || other.ndim() < 2 {
return Err(Error::Shape(format!(
@@ -153,7 +132,7 @@ impl Tensor {
Ok(())
}
- /// Matching inner dimensions and compatible batches; rank is already checked.
+ /// Assumes the rank check already passed.
fn check_matmul_dims(&self, other: &Tensor, inner: usize, other_inner: usize) -> Result<()> {
if inner != other_inner {
return Err(Error::Shape(format!(
diff --git a/src/tensor/core.rs b/src/tensor/core.rs
index 0755a94..cfb5bc2 100644
--- a/src/tensor/core.rs
+++ b/src/tensor/core.rs
@@ -8,16 +8,10 @@ use super::shape;
use super::storage::Storage;
use super::Device;
-/// An n-dimensional array of `f32`, on the CPU or a GPU.
+/// An n-dimensional `f32` array on the CPU or a GPU.
///
-/// Tensors are immutable values. Cloning one bumps a refcount rather than
-/// copying data, so passing them around is cheap and ops return new tensors
-/// instead of mutating in place.
-///
-/// A tensor either carries a graph node — meaning it descends from a parameter
-/// and gradients flow through it — or it does not, in which case it is plain
-/// data. There is no `requires_grad` flag to keep in sync: see
-/// [`grad_fn`](Tensor::grad_fn).
+/// Immutable, and cloning only bumps a refcount. A tensor with a graph node
+/// takes part in autograd, one without is plain data (see [`grad_fn`](Tensor::grad_fn)).
#[derive(Clone)]
pub struct Tensor {
storage: Arc,
@@ -27,8 +21,6 @@ pub struct Tensor {
}
impl Tensor {
- // ── Construction ─────────────────────────────────────────────────────────
-
/// Wrap `data` in a tensor of `shape` on the CPU.
///
/// Panics if `data.len()` does not match the shape.
@@ -47,14 +39,12 @@ impl Tensor {
Tensor::from_vec(vec![value], &[1])
}
- /// Build a tensor directly from storage. Used by ops that produce their
- /// output on the GPU without a host round trip.
+ /// Wrap storage an op already produced, possibly on the GPU.
pub(crate) fn raw(storage: Storage, shape: Vec, device: Device) -> Tensor {
Tensor::raw_shared(Arc::new(storage), shape, device)
}
- /// Reuse an existing allocation under a new shape. This is what makes
- /// `reshape` free — the tensors alias the same buffer.
+ /// Reuse an allocation under a new shape, which is what makes `reshape` free.
pub(crate) fn raw_shared(storage: Arc, shape: Vec, device: Device) -> Tensor {
Tensor {
storage,
@@ -79,8 +69,6 @@ impl Tensor {
}
}
- // ── Shape and metadata ───────────────────────────────────────────────────
-
pub fn shape(&self) -> &[usize] {
&self.shape
}
@@ -111,14 +99,12 @@ impl Tensor {
self.shape[axis]
}
- /// The last dimension — the feature axis for most layers.
+ /// The last dimension.
pub fn last_dim(&self) -> usize {
*self.shape.last().expect("tensor has no dimensions")
}
- // ── Data access ──────────────────────────────────────────────────────────
-
- /// A host copy of the data. Downloads from the GPU when needed.
+ /// A host copy of the data.
pub fn to_vec(&self) -> Vec {
self.storage.to_vec()
}
@@ -147,8 +133,6 @@ impl Tensor {
self.storage.clone()
}
- // ── Devices ──────────────────────────────────────────────────────────────
-
/// A copy of this tensor on `device`, differentiable in both directions.
pub fn to(&self, device: Device) -> Tensor {
if self.device == device {
@@ -164,18 +148,13 @@ impl Tensor {
self.to(Device::Cpu)
}
- /// Shorthand for `to(Device::Cuda(0))`. Panics if there is no GPU — use
- /// [`Device::cuda`] when you want to handle that.
+ /// Shorthand for `to(Device::Cuda(0))`. Panics without a GPU, see [`Device::cuda`].
pub fn cuda(&self) -> Tensor {
self.to(Device::cuda(0).expect("cuda: no device"))
}
- // ── Autograd ─────────────────────────────────────────────────────────────
-
- /// The graph node that produced this tensor, if any.
- ///
- /// `None` means plain data: a literal, a dataset batch, or anything built
- /// inside [`no_grad`](crate::autograd::no_grad).
+ /// The graph node that produced this tensor. `None` for plain data, such as
+ /// a literal or anything built inside [`no_grad`](crate::autograd::no_grad).
pub fn grad_fn(&self) -> Option<&Node> {
self.node.as_ref()
}
@@ -195,11 +174,9 @@ impl Tensor {
autograd::backward(self);
}
- /// Attach `rule` as this tensor's backward, if any input is differentiable
- /// and tracking is on.
- ///
- /// `rule` is a closure so that saving forward values costs nothing during
- /// inference. Its gradients must come back in the same order as `inputs`.
+ /// Attach `rule` as the backward if tracking is on and any input needs it.
+ /// It is a closure so inference never pays for saved values. Gradients must
+ /// come back in the order of `inputs`.
pub(crate) fn with_grad(mut self, inputs: &[&Tensor], rule: impl FnOnce() -> R) -> Tensor
where
R: Backward + 'static,
@@ -212,7 +189,7 @@ impl Tensor {
}
}
-/// Moving between devices is a copy in both directions; the gradient just travels back.
+/// The gradient of a device copy is a copy back.
struct ToDevice {
from: Device,
}
diff --git a/src/tensor/device.rs b/src/tensor/device.rs
index fbe6bca..69b9d9d 100644
--- a/src/tensor/device.rs
+++ b/src/tensor/device.rs
@@ -1,13 +1,11 @@
-//! Where a tensor's data lives.
+//! Where a tensor lives.
use std::fmt;
use crate::error::{Error, Result};
-/// The device a tensor's storage is allocated on.
-///
-/// Ops require both operands on the same device; mixing them panics. Move data
-/// explicitly with [`Tensor::to`](crate::tensor::Tensor::to).
+/// Where a tensor's storage is allocated. Ops panic if operands disagree, so
+/// move data with [`Tensor::to`](crate::tensor::Tensor::to).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Device {
Cpu,
diff --git a/src/tensor/display.rs b/src/tensor/display.rs
index 88b12de..3a8c7c8 100644
--- a/src/tensor/display.rs
+++ b/src/tensor/display.rs
@@ -1,15 +1,11 @@
-//! Printing tensors.
-//!
-//! `{:?}` gives the one-line summary you want in a log; `{}` prints values, with
-//! long runs elided so a debug print never floods a terminal.
+//! `{:?}` is a one-line summary, `{}` prints the values (long runs elided).
use std::fmt;
use super::Tensor;
-/// Values shown at each end before eliding the middle.
+/// Values kept at each end of an elided run.
const EDGE: usize = 3;
-/// Rows shown before eliding the rest of a matrix.
const MAX_ROWS: usize = 6;
impl fmt::Debug for Tensor {
@@ -50,8 +46,7 @@ impl fmt::Display for Tensor {
}
write!(f, "] {:?}", self)
}
- // Higher ranks have no layout that reads well in a terminal; the
- // summary plus the first few values is more useful than a wall of text.
+ // no layout reads well past 2-D, so print the flat edges
_ => {
write!(f, "[")?;
write_row(f, &data)?;
diff --git a/src/tensor/init.rs b/src/tensor/init.rs
index 3654668..ecea92e 100644
--- a/src/tensor/init.rs
+++ b/src/tensor/init.rs
@@ -1,7 +1,4 @@
-//! Ways to make a tensor full of values.
-//!
-//! Everything here builds on the CPU. Chain `.to(device)` to place the result,
-//! or let a layer's `to_device` move its parameters after construction.
+//! Tensor constructors. They build on the CPU, so chain `.to(device)` to place the result.
use crate::rng;
@@ -44,17 +41,13 @@ impl Tensor {
Tensor::from_vec(rng::uniform(shape::numel(shape), lo, hi), shape)
}
- /// Kaiming (He) uniform — the default for ReLU-family layers.
- ///
- /// Bounds are `±√(6 / fan_in)`, which keeps activation variance roughly
- /// constant through a stack of rectified layers.
+ /// Kaiming (He) uniform, bounds `±√(6 / fan_in)`. For ReLU-family layers.
pub fn kaiming_uniform(shape: &[usize], fan_in: usize) -> Tensor {
let bound = (6.0 / fan_in as f32).sqrt();
Tensor::uniform(shape, -bound, bound)
}
- /// Xavier (Glorot) uniform — for `tanh`/`sigmoid` layers, where the backward
- /// variance matters as much as the forward.
+ /// Xavier (Glorot) uniform, for `tanh` and `sigmoid` layers.
pub fn xavier_uniform(shape: &[usize], fan_in: usize, fan_out: usize) -> Tensor {
let bound = (6.0 / (fan_in + fan_out) as f32).sqrt();
Tensor::uniform(shape, -bound, bound)
@@ -88,7 +81,7 @@ impl Tensor {
Tensor::from_vec(data, &[n, n])
}
- /// Rows of the identity, selected by `indices` — one-hot targets in `classes` classes.
+ /// One-hot rows for `indices`, `[indices.len(), classes]`.
pub fn one_hot(indices: &[usize], classes: usize) -> Tensor {
let mut data = vec![0.0f32; indices.len() * classes];
for (row, &class) in indices.iter().enumerate() {
diff --git a/src/tensor/mod.rs b/src/tensor/mod.rs
index bf91db5..88de50d 100644
--- a/src/tensor/mod.rs
+++ b/src/tensor/mod.rs
@@ -1,18 +1,4 @@
-//! Tensors: the data structure everything else is built on.
-//!
-//! - [`Tensor`] — shape, storage, and its link into the autograd graph (`core.rs`)
-//! - [`Device`] — CPU or a numbered GPU
-//! - [`ops`] — every operation, grouped by what it does
-//! - [`shape`] — strides, broadcasting, and index arithmetic
-//! - [`storage`] — the bytes, on one device or the other
-//!
-//! ```
-//! use fastnn::prelude::*;
-//!
-//! let x = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]);
-//! let y = x.matmul(&x).relu().sum();
-//! assert_eq!(y.item(), 54.0); // [[7, 10], [15, 22]] summed
-//! ```
+//! Tensors: shape, storage and the link into the autograd graph.
mod checked;
mod core;
diff --git a/src/tensor/ops/activation.rs b/src/tensor/ops/activation.rs
index 11fe466..66aea76 100644
--- a/src/tensor/ops/activation.rs
+++ b/src/tensor/ops/activation.rs
@@ -1,7 +1,5 @@
-//! Activation functions and the softmax family.
-//!
-//! Softmax normalises over the **last** dimension. Anything with a different
-//! layout should `reshape` to `[rows, classes]` first — attention does exactly that.
+//! Activations and the softmax family. Softmax runs over the last dimension, so
+//! reshape to `[rows, classes]` first for other layouts.
use crate::autograd::ops::activation::{
GeluBackward, LeakyReluBackward, LogSoftmaxBackward, ReluBackward, SigmoidBackward,
@@ -35,7 +33,7 @@ impl Tensor {
out.with_grad(&[self], || TanhBackward { output: saved })
}
- /// Exact GELU, `x · Φ(x)` — the transformer default.
+ /// Exact GELU, `x · Φ(x)`, the transformer default.
pub fn gelu(&self) -> Tensor {
unary_op(self, gelu_scalar, ffi::fastnn_cuda_gelu).with_grad(&[self], || GeluBackward {
input: self.detach(),
@@ -94,8 +92,7 @@ impl Tensor {
out.with_grad(&[self], || SoftmaxBackward { output: saved })
}
- /// `log(softmax(x))` over the last dimension, computed without the
- /// intermediate exponential so large logits do not overflow.
+ /// `log(softmax(x))` over the last dimension, without the overflow-prone exp.
pub fn log_softmax(&self) -> Tensor {
let (rows, cols) = self.as_rows();
let out = match self.storage() {
@@ -120,7 +117,7 @@ impl Tensor {
out.with_grad(&[self], || LogSoftmaxBackward { output: saved })
}
- /// View the tensor as `(rows, last_dim)` — the layout every row-wise op uses.
+ /// The tensor as `(rows, last_dim)`.
pub(crate) fn as_rows(&self) -> (usize, usize) {
let cols = self.last_dim();
(self.numel() / cols, cols)
@@ -135,7 +132,7 @@ fn gelu_scalar(x: f32) -> f32 {
0.5 * x * (1.0 + libm::erff(x * std::f32::consts::FRAC_1_SQRT_2))
}
-/// Largest value in a row; subtracting it keeps `exp` in range.
+// subtracted before exp so large logits do not overflow
fn row_max(row: &[f32]) -> f32 {
row.iter().copied().fold(f32::NEG_INFINITY, f32::max)
}
diff --git a/src/tensor/ops/arith.rs b/src/tensor/ops/arith.rs
index 03e81ce..6c30ed5 100644
--- a/src/tensor/ops/arith.rs
+++ b/src/tensor/ops/arith.rs
@@ -1,4 +1,4 @@
-//! Element-wise arithmetic, with NumPy broadcasting.
+//! Element-wise arithmetic with broadcasting.
use crate::autograd::ops::arith::{
AddBackward, DivBackward, MulBackward, MulScalarBackward, SubBackward,
@@ -66,8 +66,7 @@ impl Tensor {
}
}
-// Operator sugar for the common infix forms. Taking references keeps the
-// no-copy semantics visible at the call site: `&a + &b`.
+// Operators are on references (`&a + &b`) so nothing gets moved.
macro_rules! impl_binary_operator {
($trait:ident, $method:ident, $op:ident) => {
impl std::ops::$trait for &Tensor {
diff --git a/src/tensor/ops/conv.rs b/src/tensor/ops/conv.rs
index 73d9697..3beedba 100644
--- a/src/tensor/ops/conv.rs
+++ b/src/tensor/ops/conv.rs
@@ -1,21 +1,12 @@
-//! `im2col` and `col2im`, the pair that turns convolution into matrix algebra.
-//!
-//! Unfolding each sliding window into a column lets [`Conv2d`](crate::nn::Conv2d)
-//! be `weight · columns`, so it inherits cuBLAS on the GPU and a ready-made
-//! derivative from [`matmul`](super::matmul). `col2im` is the exact adjoint —
-//! the transpose of the same linear map — which makes it both `im2col`'s
-//! gradient and the forward pass of transposed convolution: the two operations
-//! are each other's backward rule.
+//! `im2col` and `col2im`, which turn convolution into a matmul. They are exact
+//! adjoints, so each is the other's backward rule.
use crate::autograd::ops::conv::{Col2ImBackward, Im2ColBackward};
use crate::cuda::kernels;
use crate::tensor::storage::Storage;
use crate::tensor::Tensor;
-/// The numbers that describe a 2-D sliding window.
-///
-/// With dilation `d`, kernel tap `k` reads offset `k·d` — the kernel's
-/// footprint stretches to `(k−1)·d + 1` pixels without adding weights.
+/// A 2-D sliding window. Dilation `d` spreads the taps to `(k-1)·d + 1` pixels.
#[derive(Clone, Copy, Debug)]
pub struct Window {
pub kernel: (usize, usize),
@@ -25,7 +16,7 @@ pub struct Window {
}
impl Window {
- /// A square window with the same value on both axes.
+ /// Same kernel, stride and padding on both axes.
pub fn square(kernel: usize, stride: usize, padding: usize) -> Window {
Window {
kernel: (kernel, kernel),
@@ -42,7 +33,6 @@ impl Window {
self
}
- /// How many pixels the kernel spans along one axis: `(k−1)·d + 1`.
fn span(kernel: usize, dilation: usize) -> usize {
(kernel - 1) * dilation + 1
}
@@ -75,8 +65,7 @@ impl Window {
)
}
- /// Source pixel for output position `out` and kernel tap `k`, or `None`
- /// if it falls in the padding.
+ /// Source pixel for output `out` and tap `k`, `None` in the padding.
fn source(
out: usize,
k: usize,
@@ -90,11 +79,8 @@ impl Window {
}
}
-/// Unfold `[N, C, H, W]` image data into `[N, C·kh·kw, positions]` columns.
-///
-/// Shared by the `im2col` forward and the `col2im` backward — the adjoint pair
-/// must use the same index map, and having it once makes that a fact rather
-/// than a discipline.
+/// `[N, C, H, W]` to `[N, C·kh·kw, positions]`. `fold` is its adjoint, and both
+/// share `Window::source` so the index maps cannot drift apart.
pub(crate) fn unfold(src: &[f32], shape: &[usize], window: Window) -> Vec {
let (n, c, h, w) = (shape[0], shape[1], shape[2], shape[3]);
let (kh, kw) = window.kernel;
@@ -141,9 +127,7 @@ pub(crate) fn unfold(src: &[f32], shape: &[usize], window: Window) -> Vec {
cols
}
-/// Fold `[N, C·kh·kw, positions]` columns back into `[N, C, H, W]`, *adding*
-/// where windows overlapped — a pixel copied into several columns gets the sum
-/// of their values back. The other half of the adjoint pair.
+/// Columns back to an image, adding where windows overlap.
pub(crate) fn fold(cols: &[f32], shape: &[usize], window: Window) -> Vec {
let (n, c, h, w) = (shape[0], shape[1], shape[2], shape[3]);
let (kh, kw) = window.kernel;
@@ -191,11 +175,7 @@ pub(crate) fn fold(cols: &[f32], shape: &[usize], window: Window) -> Vec {
}
impl Tensor {
- /// Unfold `[N, C, H, W]` into `[N, C·kh·kw, out_h·out_w]`.
- ///
- /// Column `l` holds every input pixel that feeds output pixel `l`, so a
- /// `[out_channels, C·kh·kw]` weight matrix times these columns is the
- /// convolution.
+ /// Unfold `[N, C, H, W]` into `[N, C·kh·kw, out_h·out_w]`, so conv is a weight matmul.
pub fn im2col(&self, window: Window) -> Tensor {
assert_eq!(
self.ndim(),
@@ -241,12 +221,8 @@ impl Tensor {
})
}
- /// Fold `[N, C·kh·kw, positions]` columns into a `[N, C, height, width]`
- /// image, summing where windows overlap — the adjoint of [`im2col`](Self::im2col).
- ///
- /// This is the forward pass of transposed convolution: what `im2col`
- /// gathers, `col2im` scatters back. Its gradient is therefore `im2col`
- /// itself.
+ /// Fold columns into a `[N, C, height, width]` image, summing overlaps. The
+ /// adjoint of [`im2col`](Self::im2col) and the forward of transposed conv.
pub fn col2im(&self, window: Window, size: (usize, usize)) -> Tensor {
assert_eq!(
self.ndim(),
diff --git a/src/tensor/ops/index.rs b/src/tensor/ops/index.rs
index b6bad10..5822713 100644
--- a/src/tensor/ops/index.rs
+++ b/src/tensor/ops/index.rs
@@ -1,7 +1,4 @@
-//! Selecting rows by index.
-//!
-//! This is the whole of an embedding lookup: gather rows forward, scatter-add
-//! the gradient back into the table.
+//! Row selection, the lookup behind embeddings.
use crate::autograd::ops::index::IndexSelectBackward;
use crate::cuda::kernels;
@@ -21,9 +18,7 @@ pub(crate) fn index_from_f32(value: f32) -> usize {
}
impl Tensor {
- /// Gather rows `ids` from this `[rows, cols]` table into `[ids.len(), cols]`.
- ///
- /// Repeated ids are fine — their gradients add up.
+ /// Gather rows `ids` from a `[rows, cols]` table. Repeated ids are fine.
pub fn index_select(&self, ids: &[usize]) -> Tensor {
assert_eq!(
self.ndim(),
diff --git a/src/tensor/ops/matmul.rs b/src/tensor/ops/matmul.rs
index 9d0ff9a..8a53f14 100644
--- a/src/tensor/ops/matmul.rs
+++ b/src/tensor/ops/matmul.rs
@@ -1,9 +1,6 @@
-//! Matrix multiplication, plain and transpose-fused.
-//!
-//! `matmul_nt` and `matmul_tn` read an operand transposed instead of building a
-//! transposed copy first. That matters because the backward pass of every linear
-//! layer needs exactly those two forms, and on the GPU they map onto a cuBLAS
-//! transpose flag — no staging buffer, no extra kernel.
+//! Matrix multiplication, plain and transpose-fused. `matmul_nt` and `matmul_tn`
+//! read an operand transposed in place, which is what a linear layer's backward
+//! needs and maps onto a cuBLAS flag.
use rayon::prelude::*;
@@ -12,7 +9,9 @@ use crate::cuda::kernels;
use crate::tensor::storage::Storage;
use crate::tensor::Tensor;
-/// Which operands are read transposed.
+/// Below this many multiply-adds the thread pool is not worth waking.
+const SERIAL_FLOPS: usize = 16_384;
+
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub(crate) enum Layout {
/// `C = A · B`, with `A` as `[m, k]` and `B` as `[k, n]`.
@@ -24,9 +23,8 @@ pub(crate) enum Layout {
}
impl Tensor {
- /// `self · other`. Both must be at least 2-D; leading dimensions batch.
- ///
- /// A batch dimension of 1 on either side broadcasts against the other.
+ /// `self · other`. Both must be 2-D or more, leading dims are batch dims
+ /// and a batch of 1 broadcasts.
pub fn matmul(&self, other: &Tensor) -> Tensor {
let out = multiply(self, other, Layout::Plain);
out.with_grad(&[self, other], || MatmulBackward {
@@ -54,7 +52,6 @@ impl Tensor {
}
}
-/// The `(m, n, k, batch)` of one multiplication, plus each side's batch count.
struct Dims {
m: usize,
n: usize,
@@ -64,7 +61,7 @@ struct Dims {
b_batch: usize,
}
-/// Shared forward for all three layouts, with no graph node attached.
+/// Forward for all three layouts, no graph node.
pub(crate) fn multiply(a: &Tensor, b: &Tensor, layout: Layout) -> Tensor {
assert_eq!(
a.device(),
@@ -81,9 +78,8 @@ pub(crate) fn multiply(a: &Tensor, b: &Tensor, layout: Layout) -> Tensor {
);
let dims = resolve_dims(a, b, layout);
- // The output's batch prefix comes from whichever side actually has one:
- // more batch items wins, and at equal counts the higher rank wins, so
- // [m,k] × [1,k,n] stays 3-D instead of collapsing to a plain matrix.
+ // Batch prefix comes from the side with more batches, then the higher rank,
+ // so [m,k] x [1,k,n] stays 3-D.
let mut out_shape: Vec = a.shape()[..a.ndim() - 2].to_vec();
if dims.b_batch > dims.a_batch || (dims.b_batch == dims.a_batch && b.ndim() > a.ndim()) {
out_shape = b.shape()[..b.ndim() - 2].to_vec();
@@ -96,8 +92,7 @@ pub(crate) fn multiply(a: &Tensor, b: &Tensor, layout: Layout) -> Tensor {
let out = cuda_gemm(x, y, &dims, layout);
Tensor::raw(Storage::Cuda(out), out_shape, a.device())
}
- // cuBLAS strided-batched needs matching batch counts; broadcast one side
- // by materialising it rather than adding a second GPU path.
+ // strided-batched cuBLAS needs equal batch counts
(Storage::Cuda(_), Storage::Cuda(_)) => {
let (x, y) = match_batches(a, b, &dims);
multiply(&x, &y, layout)
@@ -114,7 +109,6 @@ fn resolve_dims(a: &Tensor, b: &Tensor, layout: Layout) -> Dims {
let (ar, ac) = (a.dim(a.ndim() - 2), a.dim(a.ndim() - 1));
let (br, bc) = (b.dim(b.ndim() - 2), b.dim(b.ndim() - 1));
- // Read each side's (rows, inner) according to which operand is transposed.
let (m, k) = match layout {
Layout::LhsT => (ac, ar),
_ => (ar, ac),
@@ -148,7 +142,7 @@ fn resolve_dims(a: &Tensor, b: &Tensor, layout: Layout) -> Dims {
}
}
-/// Tile whichever side has a single batch up to the other's batch count.
+/// Tile the single-batch side up to the other's batch count.
fn match_batches(a: &Tensor, b: &Tensor, dims: &Dims) -> (Tensor, Tensor) {
let tile = |t: &Tensor| {
let mut shape = vec![dims.batch];
@@ -186,16 +180,8 @@ fn cuda_gemm(
result.expect("cuda gemm")
}
-/// Row-parallel GEMM.
-///
-/// Parallelism is per *output row*, not per batch: a transformer step's largest
-/// multiplies — the feed-forward and head projections — have batch 1, and
-/// batch-level chunks would leave every core but one idle on exactly the work
-/// that dominates.
-///
-/// Each layout gets the loop order that streams its operands contiguously:
-/// `Plain` and `LhsT` accumulate whole rows of `b` at a time, and `RhsT` — the
-/// `q·kᵀ` of every attention score — is a dot product of two contiguous rows.
+/// Parallel over output rows rather than batches, since the big transformer
+/// matmuls have batch 1. Loop order keeps each layout reading contiguously.
fn cpu_gemm(a: &[f32], b: &[f32], dims: &Dims, layout: Layout) -> Vec {
let Dims {
m,
@@ -239,8 +225,7 @@ fn cpu_gemm(a: &[f32], b: &[f32], dims: &Dims, layout: Layout) -> Vec {
};
let mut out = vec![0.0f32; batch * m * n];
- // Tiny multiplies are not worth a trip through the thread pool.
- if batch * m * n * k < 16_384 {
+ if batch * m * n * k < SERIAL_FLOPS {
out.chunks_mut(n)
.enumerate()
.for_each(|(r, row)| fill_row(r, row));
diff --git a/src/tensor/ops/mod.rs b/src/tensor/ops/mod.rs
index 6c46a7b..7972599 100644
--- a/src/tensor/ops/mod.rs
+++ b/src/tensor/ops/mod.rs
@@ -1,22 +1,5 @@
-//! Tensor operations, grouped by what they do.
-//!
-//! | module | ops |
-//! |---|---|
-//! | [`arith`] | `add` `sub` `mul` `div` and their scalar forms |
-//! | [`unary`] | `neg` `exp` `log` `sqrt` `abs` `powf` `clamp` |
-//! | [`activation`] | `relu` `sigmoid` `tanh` `gelu` `silu` `leaky_relu` `softmax` `log_softmax` |
-//! | [`matmul`] | `matmul` and the transpose-fused `matmul_nt` / `matmul_tn` |
-//! | [`reduce`] | `sum` `mean` `max` `min` `argmax` and their axis forms |
-//! | [`view`] | `reshape` `permute` `transpose` `expand` `cat` `stack` `narrow` |
-//! | [`index`] | `index_select` |
-//! | [`conv`] | `im2col`, the unfold that turns convolution into a GEMM |
-//! | [`pool`] | `max_pool2d` `avg_pool2d` `adaptive_avg_pool2d` |
-//! | [`norm`] | `layer_norm` `batch_norm2d`, fused so the backward can reuse the statistics |
-//!
-//! Each public op computes its result on whichever device the inputs live on,
-//! then attaches a backward rule from [`autograd::ops`](crate::autograd::ops).
-//! The `raw_*` helpers here do the same compute with no graph, for use inside
-//! backward rules where recording would be wrong.
+//! Tensor ops, one module per family. Each computes on the inputs' device and
+//! attaches a rule from [`autograd::ops`](crate::autograd::ops).
pub mod activation;
pub mod arith;
@@ -37,10 +20,9 @@ use super::shape;
use super::storage::Storage;
use super::Tensor;
-/// Elements below which rayon's overhead outweighs the parallelism.
+/// Below this many elements rayon costs more than it saves.
const PARALLEL_THRESHOLD: usize = 1 << 15;
-/// Apply an element-wise function of one tensor.
pub(crate) fn unary_op(
x: &Tensor,
cpu: impl Fn(f32) -> f32 + Send + Sync,
@@ -55,7 +37,6 @@ pub(crate) fn unary_op(
}
}
-/// Apply an element-wise function of a tensor and a scalar.
pub(crate) fn scalar_op(
x: &Tensor,
s: f32,
@@ -71,7 +52,6 @@ pub(crate) fn scalar_op(
}
}
-/// Apply an element-wise function of two tensors, broadcasting if their shapes differ.
pub(crate) fn binary_op(
a: &Tensor,
b: &Tensor,
@@ -81,7 +61,7 @@ pub(crate) fn binary_op(
assert_eq!(
a.device(),
b.device(),
- "device mismatch: {} and {} — move one with .to()",
+ "device mismatch: {} and {}; move one with .to()",
a.device(),
b.device()
);
@@ -100,10 +80,7 @@ pub(crate) fn binary_op(
}
}
-/// Broadcast both operands to a common shape, then apply `cpu`/`gpu`.
-///
-/// On the GPU we materialise the expansion so the kernel stays on-device; on the
-/// CPU we walk broadcast strides instead, which avoids the copy entirely.
+/// GPU expands both sides first, CPU walks broadcast strides and skips the copy.
fn broadcast_binary_op(
a: &Tensor,
b: &Tensor,
@@ -140,7 +117,6 @@ fn broadcast_binary_op(
Tensor::from_vec(data, &out_shape)
}
-/// Map over host data, in parallel once it is large enough to pay off.
pub(crate) fn map(data: &[f32], f: impl Fn(f32) -> f32 + Send + Sync) -> Vec {
if data.len() >= PARALLEL_THRESHOLD {
data.par_iter().map(|&v| f(v)).collect()
@@ -149,7 +125,6 @@ pub(crate) fn map(data: &[f32], f: impl Fn(f32) -> f32 + Send + Sync) -> Vec f32 + Send + Sync) -> Vec {
debug_assert_eq!(a.len(), b.len());
if a.len() >= PARALLEL_THRESHOLD {
diff --git a/src/tensor/ops/norm.rs b/src/tensor/ops/norm.rs
index fc7cf8e..318e70f 100644
--- a/src/tensor/ops/norm.rs
+++ b/src/tensor/ops/norm.rs
@@ -1,10 +1,5 @@
-//! Fused normalization ops.
-//!
-//! Both of these could be assembled from `mean`, `sub`, `sqrt`, and `div` — and
-//! [`RMSNorm`](crate::nn::RMSNorm) is. They are fused instead because the
-//! statistics are needed again in the backward pass, and recomputing them
-//! through a chain of eight small ops costs more memory and more kernels than
-//! saving two vectors.
+//! Fused layer and batch norm. They could be built from small ops but the
+//! backward needs the statistics again, and saving them is cheaper than the graph.
use crate::autograd::ops::norm::{BatchNorm2dBackward, LayerNormBackward, LayerNormCudaBackward};
use crate::cuda::kernels;
@@ -71,9 +66,8 @@ impl Tensor {
})
}
- /// Normalize `[N, C, H, W]` per channel using this batch's own statistics,
- /// then scale and shift. Training-mode only — at eval time the layer applies
- /// its running statistics as a plain affine transform.
+ /// Normalize `[N, C, H, W]` per channel with this batch's statistics. Training
+ /// only, eval uses the layer's running statistics instead.
pub fn batch_norm2d(&self, gamma: &Tensor, beta: &Tensor, eps: f32) -> (Tensor, BatchStats) {
assert_eq!(
self.ndim(),
diff --git a/src/tensor/ops/pool.rs b/src/tensor/ops/pool.rs
index 9c5d84c..6a0ace8 100644
--- a/src/tensor/ops/pool.rs
+++ b/src/tensor/ops/pool.rs
@@ -1,7 +1,4 @@
-//! 2-D pooling over `[N, C, H, W]` tensors.
-//!
-//! Pooling is a reduction with a routing decision, so unlike convolution it does
-//! not reduce to a matrix multiply — each variant carries its own gradient rule.
+//! 2-D pooling over `[N, C, H, W]`. Each variant has its own gradient rule.
use std::ops::Range;
@@ -11,10 +8,7 @@ use crate::tensor::Tensor;
use super::conv::Window;
impl Tensor {
- /// Take the largest value in each window.
- ///
- /// The backward pass needs to know *which* input won, so the forward saves
- /// one flat index per output element.
+ /// Largest value in each window. Saves the winning index for the backward.
pub fn max_pool2d(&self, window: Window) -> Tensor {
let (n, c, h, w) = self.expect_nchw("max_pool2d");
let (out_h, out_w) = window.output_size(h, w);
@@ -47,7 +41,7 @@ impl Tensor {
.with_grad(&[self], || MaxPool2dBackward { shape, winners })
}
- /// Average the values in each window, counting only real pixels (not padding).
+ /// Average of each window, padding not counted.
pub fn avg_pool2d(&self, window: Window) -> Tensor {
let (n, c, h, w) = self.expect_nchw("avg_pool2d");
let (out_h, out_w) = window.output_size(h, w);
@@ -74,10 +68,7 @@ impl Tensor {
.with_grad(&[self], || AvgPool2dBackward { shape, window })
}
- /// Average into a fixed `(height, width)` output, whatever the input size.
- ///
- /// `(1, 1)` is global average pooling — the usual bridge from a conv stack
- /// to a classifier head.
+ /// Average into a fixed `(height, width)` output. `(1, 1)` is global pooling.
pub fn adaptive_avg_pool2d(&self, output: (usize, usize)) -> Tensor {
let (n, c, h, w) = self.expect_nchw("adaptive_avg_pool2d");
let (out_h, out_w) = output;
diff --git a/src/tensor/ops/reduce.rs b/src/tensor/ops/reduce.rs
index 8f68325..ab0315c 100644
--- a/src/tensor/ops/reduce.rs
+++ b/src/tensor/ops/reduce.rs
@@ -1,8 +1,5 @@
-//! Reductions over all elements or along one axis.
-//!
-//! `sum` and `mean` are differentiable. `max`, `min`, and `argmax` return plain
-//! numbers rather than tensors — they are selection, not arithmetic, and giving
-//! them a tensor type would suggest a gradient they do not have.
+//! Reductions over everything or along one axis. `max`, `min` and `argmax`
+//! return plain numbers, since they are selections with no gradient.
use crate::autograd::ops::reduce::{MeanBackward, SumAxisBackward, SumBackward};
use crate::cuda::kernels;
diff --git a/src/tensor/ops/unary.rs b/src/tensor/ops/unary.rs
index c9837d1..1ab37fe 100644
--- a/src/tensor/ops/unary.rs
+++ b/src/tensor/ops/unary.rs
@@ -1,4 +1,4 @@
-//! Element-wise maths on a single tensor.
+//! Element-wise maths on one tensor.
use crate::autograd::ops::unary::{
AbsBackward, ClampBackward, ExpBackward, LogBackward, NegBackward, PowBackward, SqrtBackward,
diff --git a/src/tensor/ops/view.rs b/src/tensor/ops/view.rs
index 44dd5b5..a5c10e8 100644
--- a/src/tensor/ops/view.rs
+++ b/src/tensor/ops/view.rs
@@ -1,8 +1,5 @@
-//! Ops that rearrange elements without doing arithmetic.
-//!
-//! Tensors are always contiguous, so `reshape` is free (it shares storage and
-//! only changes the shape vector) while `permute`, `expand`, and `cat` write a
-//! new buffer.
+//! Ops that rearrange elements without arithmetic. `reshape` shares storage,
+//! the rest write a new buffer.
use crate::autograd::ops::view::{
CatBackward, ExpandBackward, NarrowBackward, PermuteBackward, ReshapeBackward,
@@ -13,9 +10,7 @@ use crate::tensor::storage::Storage;
use crate::tensor::Tensor;
impl Tensor {
- /// Reinterpret the elements under a new shape. One dimension may be `-1`.
- ///
- /// Free: the result shares storage with the input.
+ /// Same elements under a new shape, sharing storage. One dim may be `-1`.
pub fn reshape(&self, dims: &[i64]) -> Tensor {
let new_shape = shape::resolve(dims, self.numel());
let old_shape = self.shape().to_vec();
@@ -27,9 +22,7 @@ impl Tensor {
self.reshape(&[-1])
}
- /// Collapse dimensions `from..` into one, keeping the dimensions before it.
- ///
- /// `flatten_from(1)` is the usual conv-to-linear bridge: `[N, C, H, W] → [N, C*H*W]`.
+ /// Collapse dimensions `from..` into one, e.g. `[N, C, H, W]` to `[N, C*H*W]` at 1.
pub fn flatten_from(&self, from: usize) -> Tensor {
let mut dims: Vec = self.shape()[..from].iter().map(|&d| d as i64).collect();
dims.push(-1);
@@ -188,7 +181,6 @@ impl Tensor {
}
}
- // Interleave: for each outer slice, append every part's chunk in order.
let (outer, _, inner) = shape::split_at_axis(&out_shape, axis);
let sources: Vec> = parts.iter().map(|t| t.to_vec()).collect();
let mut data = Vec::with_capacity(shape::numel(&out_shape));
@@ -213,8 +205,6 @@ impl Tensor {
}
}
-// ── Shared machinery, also used by the backward rules ────────────────────────
-
/// Same shape change as [`Tensor::reshape`], with no graph node.
pub(crate) fn reshaped(x: &Tensor, shape: Vec) -> Tensor {
debug_assert_eq!(shape::numel(&shape), x.numel());
@@ -226,8 +216,7 @@ pub(crate) fn expanded(x: &Tensor, to: &[usize]) -> Tensor {
if x.shape() == to {
return x.detach();
}
- // Stride 0 on a broadcast dimension makes every output index read slot 0,
- // which is exactly what the permute kernel needs to double as `expand`.
+ // stride 0 on stretched dims lets the permute gather double as expand
let in_strides: Vec = shape::strides_for(x.shape())
.iter()
.enumerate()
@@ -246,10 +235,7 @@ pub(crate) fn broadcast_to(x: &Tensor, to: &[usize]) -> Tensor {
expanded(&reshaped(x, padded), to)
}
-/// Build `out_shape` by reading each output position from `in_strides[order[d]]`.
-///
-/// This one gather covers both `permute` (a real reordering) and `expand`
-/// (identity order, zero strides on broadcast dimensions).
+/// Read each output position through `in_strides[order[d]]`. Serves both permute and expand.
fn gather(x: &Tensor, out_shape: Vec, in_strides: &[usize], order: &[usize]) -> Tensor {
let numel = shape::numel(&out_shape);
let out_strides = shape::strides_for(&out_shape);
diff --git a/src/tensor/shape.rs b/src/tensor/shape.rs
index 8f35f56..1ab7ca3 100644
--- a/src/tensor/shape.rs
+++ b/src/tensor/shape.rs
@@ -1,8 +1,5 @@
-//! Shape arithmetic: strides, broadcasting, and flat-index math.
-//!
-//! Tensors are always C-contiguous, so strides are derived from the shape rather
-//! than stored as an independent source of truth. Every view op (`reshape`,
-//! `permute`, `expand`) materialises a new contiguous buffer.
+//! Strides, broadcasting and index arithmetic. Tensors are always contiguous, so
+//! strides are derived from the shape instead of stored.
/// Row-major strides for `shape`, in elements.
pub fn strides_for(shape: &[usize]) -> Vec {
@@ -18,9 +15,7 @@ pub fn numel(shape: &[usize]) -> usize {
shape.iter().product()
}
-/// Resolve a shape that may contain a single `-1` placeholder against `numel`.
-///
-/// Panics if more than one dimension is `-1` or the element count does not divide.
+/// Fill in a single `-1` placeholder so the shape holds `numel` elements.
pub fn resolve(shape: &[i64], numel: usize) -> Vec {
let mut placeholder = None;
let mut known = 1usize;
@@ -60,9 +55,7 @@ pub fn resolve(shape: &[i64], numel: usize) -> Vec {
}
}
-/// The shape two operands broadcast to, following NumPy rules.
-///
-/// Panics if the shapes are incompatible.
+/// The shape two operands broadcast to, NumPy style.
pub fn broadcast(a: &[usize], b: &[usize]) -> Vec {
let ndim = a.len().max(b.len());
(0..ndim)
@@ -85,8 +78,7 @@ pub fn pad_left(shape: &[usize], ndim: usize) -> Vec {
out
}
-/// Strides for reading `shape` (already padded to `out.len()`) while broadcasting
-/// to `out`. Broadcast dimensions get stride 0 so every output index reads slot 0.
+/// Strides for reading the padded `shape` as `out`. Stretched dims get stride 0.
pub fn broadcast_strides(shape: &[usize], out: &[usize]) -> Vec {
debug_assert_eq!(shape.len(), out.len());
strides_for(shape)
@@ -97,8 +89,6 @@ pub fn broadcast_strides(shape: &[usize], out: &[usize]) -> Vec {
}
/// Split `shape` around `axis` into (product before, size at, product after).
-///
-/// Most reductions and axis-wise scatters reduce to a triple loop over these.
pub fn split_at_axis(shape: &[usize], axis: usize) -> (usize, usize, usize) {
assert!(
axis < shape.len(),
diff --git a/src/tensor/storage.rs b/src/tensor/storage.rs
index 34450f9..3aac068 100644
--- a/src/tensor/storage.rs
+++ b/src/tensor/storage.rs
@@ -1,8 +1,4 @@
-//! The bytes behind a tensor, on one device or the other.
-//!
-//! `Storage` is never cloned during normal use — a [`Tensor`](super::Tensor)
-//! shares it through an `Arc`, so passing tensors around is a refcount bump and
-//! a device copy only happens when you ask for one.
+//! The buffer behind a tensor, on the CPU or a GPU. Tensors share it through an `Arc`.
use crate::cuda::CudaBuffer;
diff --git a/tests/conv.rs b/tests/conv.rs
index 54c67e3..0a70886 100644
--- a/tests/conv.rs
+++ b/tests/conv.rs
@@ -1,17 +1,19 @@
-//! Convolution-variant semantics, checked against the definitions.
+//! Convolution variants checked against their definitions.
use fastnn::prelude::*;
use fastnn::tensor::Window;
+const TOLERANCE: f32 = 1e-4;
+
#[track_caller]
fn assert_close(got: f32, want: f32, what: &str) {
- assert!((got - want).abs() < 1e-4, "{what}: got {got}, want {want}");
+ assert!(
+ (got - want).abs() < TOLERANCE,
+ "{what}: got {got}, want {want}"
+ );
}
-/// Transposed convolution IS the adjoint of convolution: with a shared weight
-/// and no bias, ⟨conv(x), y⟩ = ⟨x, convᵀ(y)⟩ for every x and y. This is the
-/// defining property — if it holds for random tensors, the layer computes the
-/// transpose of the same linear map, output padding, stride, and all.
+/// With a shared weight and no bias, = .
#[test]
fn conv_transpose_is_the_exact_adjoint_of_conv() {
manual_seed(12);
@@ -19,14 +21,10 @@ fn conv_transpose_is_the_exact_adjoint_of_conv() {
let conv = Conv2d::with_window(2, 3, window, false);
let transpose = ConvTranspose2d::with_window(3, 2, window, false);
- // Same matrix in both layers. Conv flattens [out, in, kh, kw] to
- // [out, in·kh·kw]; the transpose flattens its [in_t=out, out_t=in, kh, kw]
- // to the same [out, in·kh·kw] and reads it transposed in the GEMM — so the
- // raw buffers coincide element for element.
+ // Both layers flatten the weight to the same [out, in*kh*kw] buffer.
transpose.weight.set_value(conv.weight.value());
- // 7 satisfies (size + 2p − k) % s = 0, so conv's input size is recoverable
- // exactly; a size like 6 would need output_padding to disambiguate.
+ // 7 satisfies (size + 2p - k) % s = 0, a size like 6 would need output_padding.
let x = Tensor::randn(&[2, 2, 7, 7]);
let forward = no_grad(|| conv.forward(&x)); // [2, 3, 4, 4]
let y = Tensor::randn(&[2, 3, 4, 4]);
@@ -35,13 +33,12 @@ fn conv_transpose_is_the_exact_adjoint_of_conv() {
let lhs = forward.mul(&y).sum().item();
let rhs = x.mul(&backward).sum().item();
assert!(
- (lhs - rhs).abs() / lhs.abs().max(1.0) < 1e-4,
+ (lhs - rhs).abs() / lhs.abs().max(1.0) < TOLERANCE,
"adjoint identity violated: ⟨conv x, y⟩ = {lhs} but ⟨x, convᵀ y⟩ = {rhs}"
);
}
-/// The smallest checkable upsample: a 1×1 input through a 2×2 all-ones kernel
-/// at stride 2 paints the value into every cell of a 2×2 output.
+/// A 1x1 input through a 2x2 all-ones kernel at stride 2 fills a 2x2 output.
#[test]
fn conv_transpose_upsamples_a_pixel_by_hand() {
let layer = ConvTranspose2d::with_window(1, 1, Window::square(2, 2, 0), false);
@@ -54,8 +51,7 @@ fn conv_transpose_upsamples_a_pixel_by_hand() {
assert_eq!(out.to_vec(), vec![5.0; 4]);
}
-/// Depthwise convolution with 1×1 kernels reduces to a per-channel scale —
-/// and the channels must not mix.
+/// With 1x1 kernels this is a per-channel scale and the channels must not mix.
#[test]
fn depthwise_keeps_channels_separate() {
let layer = Conv2d::depthwise(2, 1, 1, 0);
@@ -75,8 +71,7 @@ fn depthwise_keeps_channels_separate() {
);
}
-/// A dilated kernel reads every d-th input: `out[t] = a·x[t] + b·x[t+2]` for
-/// kernel `[a, b]` at dilation 2, worked by hand on a length-5 sequence.
+/// Kernel `[a, b]` at dilation 2 gives `out[t] = a*x[t] + b*x[t+2]`.
#[test]
fn dilation_stretches_the_kernel_footprint() {
let window = Window {
@@ -92,16 +87,14 @@ fn dilation_stretches_the_kernel_footprint() {
let x = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0], &[1, 1, 1, 5]);
let out = no_grad(|| layer.forward(&x)).to_vec();
- // Footprint spans 3, so 3 outputs: 10·x[t] + x[t+2].
assert_eq!(out, vec![13.0, 24.0, 35.0]);
}
-/// Causality, tested as a property: perturbing the future must not change the
-/// past. Length is preserved, and output t sees only inputs ≤ t.
+/// Poking a later input must not change earlier outputs.
#[test]
fn causal_conv1d_cannot_see_the_future() {
manual_seed(13);
- let layer = Conv1d::causal(1, 1, 3, 2); // footprint (3−1)·2+1 = 5, all left
+ let layer = Conv1d::causal(1, 1, 3, 2);
let x = Tensor::randn(&[1, 1, 8]);
let mut poked = x.to_vec();
@@ -119,8 +112,7 @@ fn causal_conv1d_cannot_see_the_future() {
assert!(a[5] != b[5], "output 5 should see input 5");
}
-/// Grouped convolution must equal running each group's slice through its own
-/// small convolution — computed here directly from the weight slices.
+/// Each group must match its own small convolution built from the weight slices.
#[test]
fn grouped_conv_equals_independent_group_convs() {
manual_seed(14);
@@ -155,7 +147,6 @@ fn grouped_conv_equals_independent_group_convs() {
}
}
-/// Gradients reach every parameter of every new layer.
#[test]
fn gradients_reach_all_conv_variant_parameters() {
let layers: Vec<(&str, Box, Tensor)> = vec![
@@ -197,8 +188,6 @@ fn gradients_reach_all_conv_variant_parameters() {
}
}
-/// `Param::set_value` on shared handles must be visible through the layer —
-/// the mechanism the hand-value tests above rely on.
#[test]
fn output_padding_hits_the_requested_size() {
let plain = ConvTranspose2d::new(1, 1, 3, 2, 1);
diff --git a/tests/convergence.rs b/tests/convergence.rs
index 74e5b0b..ca607c1 100644
--- a/tests/convergence.rs
+++ b/tests/convergence.rs
@@ -1,11 +1,28 @@
-//! End-to-end convergence: every major layer family, trained until it works.
-//!
-//! `gradcheck` proves each derivative in isolation; these prove the pieces
-//! compose into models that actually learn. Each case is deliberately tiny so
-//! the whole file runs in seconds even in a debug build.
+//! Tiny models of each layer family, trained until they fit.
use fastnn::prelude::*;
+/// Steps until the loss is under `target` or `steps` run out, returns the last loss.
+fn train_until(
+ opt: &mut dyn Optimizer,
+ steps: usize,
+ target: f32,
+ mut compute_loss: impl FnMut() -> Tensor,
+) -> f32 {
+ let mut value = f32::MAX;
+ for _ in 0..steps {
+ let loss = compute_loss();
+ opt.zero_grad();
+ loss.backward();
+ opt.step();
+ value = loss.item();
+ if value < target {
+ break;
+ }
+ }
+ value
+}
+
#[test]
fn mlp_converges_on_xor() {
manual_seed(3);
@@ -18,17 +35,9 @@ fn mlp_converges_on_xor() {
.add(Linear::new(8, 2));
let mut opt = Adam::new(model.parameters(), 5e-2);
- let mut loss_value = f32::MAX;
- for _ in 0..1000 {
- let loss = cross_entropy(&model.forward(&inputs), &targets);
- opt.zero_grad();
- loss.backward();
- opt.step();
- loss_value = loss.item();
- if loss_value < 0.01 {
- break;
- }
- }
+ let loss_value = train_until(&mut opt, 1000, 0.01, || {
+ cross_entropy(&model.forward(&inputs), &targets)
+ });
assert!(
loss_value < 0.01,
"xor did not converge: final loss {loss_value}"
@@ -38,7 +47,7 @@ fn mlp_converges_on_xor() {
#[test]
fn linear_regression_recovers_slope_and_intercept() {
manual_seed(4);
- // y = 2x + 1, no noise: the model should recover both numbers nearly exactly.
+ // y = 2x + 1 with no noise.
let xs: Vec = (0..32).map(|i| i as f32 / 16.0 - 1.0).collect();
let ys: Vec = xs.iter().map(|x| 2.0 * x + 1.0).collect();
let inputs = Tensor::from_vec(xs, &[32, 1]);
@@ -74,17 +83,9 @@ fn cnn_overfits_a_small_batch() {
.add(Linear::new(4 * 4 * 4, 10));
let mut opt = Adam::new(model.parameters(), 1e-2);
- let mut loss_value = f32::MAX;
- for _ in 0..200 {
- let loss = cross_entropy(&model.forward(&inputs), &targets);
- opt.zero_grad();
- loss.backward();
- opt.step();
- loss_value = loss.item();
- if loss_value < 0.05 {
- break;
- }
- }
+ let loss_value = train_until(&mut opt, 200, 0.05, || {
+ cross_entropy(&model.forward(&inputs), &targets)
+ });
assert!(
loss_value < 0.05,
"cnn did not overfit ten samples: loss {loss_value}"
@@ -94,7 +95,6 @@ fn cnn_overfits_a_small_batch() {
#[test]
fn lstm_memorizes_a_short_sequence() {
manual_seed(6);
- // One fixed input sequence, one fixed target per step: pure memorization.
let inputs = Tensor::randn(&[1, 6, 4]);
let targets = [2usize, 0, 3, 1, 2, 0];
@@ -104,19 +104,10 @@ fn lstm_memorizes_a_short_sequence() {
params.extend(head.parameters());
let mut opt = Adam::new(params, 1e-2);
- let mut loss_value = f32::MAX;
- for _ in 0..300 {
- let hidden = lstm.forward(&inputs);
- let logits = head.forward(&hidden).reshape(&[6, -1]);
- let loss = cross_entropy(&logits, &targets);
- opt.zero_grad();
- loss.backward();
- opt.step();
- loss_value = loss.item();
- if loss_value < 0.05 {
- break;
- }
- }
+ let loss_value = train_until(&mut opt, 300, 0.05, || {
+ let logits = head.forward(&lstm.forward(&inputs)).reshape(&[6, -1]);
+ cross_entropy(&logits, &targets)
+ });
assert!(
loss_value < 0.05,
"lstm did not memorize: loss {loss_value}"
@@ -130,21 +121,12 @@ fn transformer_overfits_a_small_batch() {
let target_in = Tensor::randn(&[2, 4, 16]);
let wanted = Tensor::randn(&[2, 4, 16]);
- // The full encoder-decoder path, driven to overfit a fixed pair.
let model = Transformer::new(16, 2, 32, 1, 1, 0.0);
let mut opt = Adam::new(model.parameters(), 3e-3);
- let mut loss_value = f32::MAX;
- for _ in 0..400 {
- let loss = mse(&model.run(&source, &target_in), &wanted);
- opt.zero_grad();
- loss.backward();
- opt.step();
- loss_value = loss.item();
- if loss_value < 0.05 {
- break;
- }
- }
+ let loss_value = train_until(&mut opt, 400, 0.05, || {
+ mse(&model.run(&source, &target_in), &wanted)
+ });
assert!(
loss_value < 0.05,
"transformer did not overfit: loss {loss_value}"
@@ -160,7 +142,7 @@ fn saved_model_predicts_identically_after_reload() {
.add(ReLU)
.add(Linear::new(12, 3));
- // A few steps first, so the weights are not just their initialization.
+ // Train a little so the weights are not just the init.
let targets = Tensor::zeros(&[4, 3]);
let mut opt = SGD::new(model.parameters(), 0.01);
for _ in 0..5 {
@@ -201,9 +183,7 @@ fn clip_grad_value_bounds_every_gradient_element() {
);
}
-/// Every optimizer, one benchmark: recover a linear map by least squares.
-/// Learning rates are per-optimizer (a shared one would test tuning, not
-/// correctness); each must reach near-zero loss on the same data.
+/// Learning rates are per optimizer, a shared one would test tuning.
#[test]
fn every_optimizer_minimizes_least_squares() {
use fastnn::nn::Param;
@@ -235,17 +215,9 @@ fn every_optimizer_minimizes_least_squares() {
let model = Linear::new(3, 1);
let mut opt = make(model.parameters());
- let mut loss_value = f32::MAX;
- for _ in 0..1500 {
- let loss = mse(&model.forward(&inputs), &targets);
- opt.zero_grad();
- loss.backward();
- opt.step();
- loss_value = loss.item();
- if loss_value < 1e-3 {
- break;
- }
- }
+ let loss_value = train_until(opt.as_mut(), 1500, 1e-3, || {
+ mse(&model.forward(&inputs), &targets)
+ });
assert!(
loss_value < 1e-3,
"{name} failed to converge: loss {loss_value}"
diff --git a/tests/cuda_parity.rs b/tests/cuda_parity.rs
index c25c1d1..a933394 100644
--- a/tests/cuda_parity.rs
+++ b/tests/cuda_parity.rs
@@ -1,17 +1,7 @@
-//! CPU/GPU equivalence.
-//!
-//! Every op with a hand-written CUDA kernel is run twice — once on each device —
-//! and the results compared. `tests/gradcheck.rs` proves the CPU maths is right;
-//! this proves the GPU agrees with it, forward and backward.
-//!
-//! cargo test --test cuda_parity
-//!
-//! Skips itself with a note when no GPU is present, so a CPU-only machine still
-//! gets a green run.
-
-// Test inputs are always passed as a slice. `&[x.clone()]` for a single input
-// costs one clone of a refcounted handle and keeps every call site reading the
-// same way as the multi-input ones.
+//! Each CUDA op run on both devices, forward and backward, and compared.
+//! Skips with a note when there is no GPU.
+
+// `&[x.clone()]` is a cheap refcount bump and keeps single and multi input calls uniform.
#![allow(clippy::cloned_ref_to_slice_refs)]
use std::sync::Mutex;
@@ -19,23 +9,21 @@ use std::sync::Mutex;
use fastnn::nn::Param;
use fastnn::prelude::*;
-/// Difference a single op may show. The kernels are built with
-/// `--use_fast_math` and reduce in a different order than the host loops, so
-/// exact equality is not on offer; anything past this is a real disagreement.
+/// The kernels use `--use_fast_math` and reduce in a different order than the host.
const OP_TOLERANCE: f32 = 2e-4;
-/// Difference a whole block may show. Dozens of ops compound their rounding, so
-/// end-to-end comparisons get more room than single ops.
+/// A whole block compounds the rounding of dozens of ops.
const MODEL_TOLERANCE: f32 = 5e-3;
+const NORM_EPS: f32 = 1e-5;
+const GAMMA: f32 = 1.1;
+const BETA: f32 = -0.2;
+
/// Tests share one GPU and one cuBLAS handle, so they run one at a time.
static GPU: Mutex<()> = Mutex::new(());
-/// Take the GPU lock, ignoring poison.
-///
-/// A failing assertion poisons the mutex, and without this every later test
-/// reports that poison instead of its own result — hiding the failure that
-/// actually matters behind a wall of noise.
+/// Ignores poison, otherwise one failed assertion makes every later test report
+/// the poison instead of its own result.
fn lock_gpu() -> std::sync::MutexGuard<'static, ()> {
GPU.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
}
@@ -45,7 +33,7 @@ fn gpu() -> Option {
match Device::cuda(0) {
Ok(device) => Some(device),
Err(_) => {
- eprintln!("no CUDA device — skipping parity test");
+ eprintln!("no CUDA device, skipping parity test");
None
}
}
@@ -80,7 +68,6 @@ fn assert_close(name: &str, cpu: &[f32], cuda: &[f32], tolerance: f32) {
}
}
-/// Run `f` on both devices and compare its output.
#[track_caller]
fn check_forward(name: &str, inputs: &[Tensor], f: impl Fn(&[Tensor]) -> Tensor) {
let Some(device) = gpu() else { return };
@@ -93,7 +80,6 @@ fn check_forward(name: &str, inputs: &[Tensor], f: impl Fn(&[Tensor]) -> Tensor)
assert_close(name, &on_cpu, &on_gpu, OP_TOLERANCE);
}
-/// Run `f` on both devices and compare every input gradient.
#[track_caller]
fn check_backward(name: &str, inputs: &[Tensor], f: impl Fn(&[Tensor]) -> Tensor) {
let Some(device) = gpu() else { return };
@@ -117,8 +103,6 @@ fn check_backward(name: &str, inputs: &[Tensor], f: impl Fn(&[Tensor]) -> Tensor
}
}
-// ── Element-wise and activations ─────────────────────────────────────────────
-
#[test]
fn elementwise() {
let a = sample(&[4, 16]);
@@ -133,8 +117,7 @@ fn elementwise() {
#[test]
fn broadcasting_stays_on_device() {
- // The GPU path materialises the expansion rather than falling back to the
- // host; the result must still match the CPU's stride walk.
+ // The GPU materialises the expansion instead of falling back to the host.
check_forward("broadcast add", &[sample(&[4, 8]), sample(&[1, 8])], |x| {
x[0].add(&x[1])
});
@@ -153,7 +136,6 @@ fn activations() {
}
}
-/// A named activation, as a plain function pointer so the list can be iterated.
type Activation = fn(&[Tensor]) -> Tensor;
fn activation_cases() -> Vec<(&'static str, Activation)> {
@@ -175,15 +157,12 @@ fn softmax() {
check_forward("softmax", &[x.clone()], |x| x[0].softmax());
check_forward("log_softmax", &[x.clone()], |x| x[0].log_softmax());
- // Weighted so the gradient is non-trivial; a plain sum over softmax is
- // constant and its gradient is zero on both devices whatever the kernel does.
+ // A plain sum over softmax has zero gradient on both devices.
check_backward("softmax", &[x, weights], |x| {
x[0].softmax().mul(&x[1]).sum()
});
}
-// ── Matrix multiplication ────────────────────────────────────────────────────
-
#[test]
fn gemm_layouts() {
check_forward("matmul", &[sample(&[8, 12]), sample(&[12, 6])], |x| {
@@ -216,7 +195,7 @@ fn gemm_batched() {
&[sample(&[4, 8, 12]), sample(&[4, 6, 12])],
|x| x[0].matmul_nt(&x[1]),
);
- // One side batched, the other shared — how a weight meets a batch of inputs.
+ // A weight meeting a batch of inputs.
check_forward(
"shared weight",
&[sample(&[8, 12]), sample(&[4, 12, 6])],
@@ -224,8 +203,6 @@ fn gemm_batched() {
);
}
-// ── Layout ───────────────────────────────────────────────────────────────────
-
#[test]
fn layout_ops() {
let x = sample(&[2, 3, 4]);
@@ -239,8 +216,6 @@ fn layout_ops() {
});
}
-// ── Reductions ───────────────────────────────────────────────────────────────
-
#[test]
fn reductions() {
let x = sample(&[6, 10]);
@@ -254,35 +229,28 @@ fn reductions() {
});
}
-// ── Fused layers ─────────────────────────────────────────────────────────────
-
#[test]
fn layer_norm() {
let x = sample(&[6, 16]);
- let gamma = Tensor::full(&[16], 1.1);
- let beta = Tensor::full(&[16], -0.2);
+ let gamma = Tensor::full(&[16], GAMMA);
+ let beta = Tensor::full(&[16], BETA);
check_forward(
"layer_norm",
&[x.clone(), gamma.clone(), beta.clone()],
- |x| x[0].layer_norm(&x[1], &x[2], 1e-5),
+ |x| x[0].layer_norm(&x[1], &x[2], NORM_EPS),
);
- // The GPU kernel produces all three gradients in one pass; the CPU path
- // computes them separately. They must agree.
- //
- // The weights come in as a fourth input rather than being built inside the
- // closure: a tensor made in there would always be on the host and would not
- // survive the GPU run.
+ // The weights are an input because a tensor built inside the closure would
+ // stay on the host.
check_backward("layer_norm", &[x, gamma, beta, sample(&[6, 16])], |x| {
- x[0].layer_norm(&x[1], &x[2], 1e-5).mul(&x[3]).sum()
+ x[0].layer_norm(&x[1], &x[2], NORM_EPS).mul(&x[3]).sum()
});
}
#[test]
fn row_reductions_handle_awkward_widths() {
- // The block reductions behind softmax and layer norm halve the thread count
- // each round, so a width that is not a power of two used to lose its tail.
- // 5 and 10 both exercise that; 10 is MNIST's class count.
+ // Non-power-of-two widths used to lose the tail of the block reduction.
+ // 10 is MNIST's class count.
for width in [3usize, 5, 7, 10, 17, 31, 100, 257] {
let x = sample(&[4, width]);
check_forward(&format!("softmax width {width}"), &[x.clone()], |x| {
@@ -293,8 +261,12 @@ fn row_reductions_handle_awkward_widths() {
});
check_forward(
&format!("layer_norm width {width}"),
- &[x, Tensor::full(&[width], 1.1), Tensor::full(&[width], -0.2)],
- |x| x[0].layer_norm(&x[1], &x[2], 1e-5),
+ &[
+ x,
+ Tensor::full(&[width], GAMMA),
+ Tensor::full(&[width], BETA),
+ ],
+ |x| x[0].layer_norm(&x[1], &x[2], NORM_EPS),
);
}
}
@@ -304,8 +276,7 @@ fn attention_matches_across_sequence_lengths() {
let Some(device) = gpu() else { return };
let _guard = lock_gpu();
- // Attention softmaxes over the key axis, so the sequence length *is* the
- // reduction width — the exact place the power-of-two bug showed up.
+ // Attention softmaxes over the key axis, so length is the reduction width.
for length in [1usize, 3, 5, 8, 13] {
manual_seed(5);
let attention = MultiHeadAttention::new(16, 4, 0.0);
@@ -336,14 +307,12 @@ fn embedding_lookup() {
check_forward("index_select", &[table.clone()], |x| {
x[0].index_select(&[0, 3, 3, 9, 1])
});
- // Row 3 twice: the GPU scatter-add must accumulate, not overwrite.
+ // Row 3 twice, the scatter-add must accumulate.
check_backward("index_select", &[table, weights], |x| {
x[0].index_select(&[0, 3, 3, 9, 1]).mul(&x[1]).sum()
});
}
-// ── End to end ───────────────────────────────────────────────────────────────
-
#[test]
fn a_transformer_block_matches() {
let Some(device) = gpu() else { return };
@@ -394,8 +363,7 @@ fn a_model_trains_identically_on_both_devices() {
#[test]
fn convolution_lowering() {
- // Stride 2 with padding 1 exercises both the padding zeros and the
- // windows-overlap accumulation in col2im.
+ // Stride 2 with padding 1 covers the padding zeros and overlap in col2im.
let x = sample(&[2, 3, 6, 5]);
let window = Window::square(3, 2, 1);
diff --git a/tests/generation.rs b/tests/generation.rs
index 9aed1a9..6253023 100644
--- a/tests/generation.rs
+++ b/tests/generation.rs
@@ -1,10 +1,8 @@
-//! The inference paths behind text generation: KV-cached decoding must match
-//! the plain forward pass, and sampling must be reproducible and honor its
-//! filters.
+//! Cached decoding against the plain forward, plus sampling and beam search.
use fastnn::prelude::*;
-/// A small causal stack in eval mode, so dropout is identity on both paths.
+/// Eval mode, so dropout is identity on both paths.
fn stack() -> TransformerStack {
manual_seed(7);
let stack = TransformerStack::causal(32, 4, 64, 2, 0.1);
@@ -12,19 +10,19 @@ fn stack() -> TransformerStack {
stack
}
+const TOLERANCE: f32 = 1e-3;
+
fn assert_close(a: &Tensor, b: &Tensor, what: &str) {
let (a, b) = (a.to_vec(), b.to_vec());
assert_eq!(a.len(), b.len(), "{what}: lengths differ");
for (i, (x, y)) in a.iter().zip(&b).enumerate() {
assert!(
- (x - y).abs() < 1e-3,
+ (x - y).abs() < TOLERANCE,
"{what}: element {i} differs, {x} vs {y}"
);
}
}
-/// Token-by-token decoding through the cache must produce the same outputs as
-/// one uncached forward over the whole sequence.
#[test]
fn cached_decoding_matches_full_forward() {
let stack = stack();
@@ -40,8 +38,7 @@ fn cached_decoding_matches_full_forward() {
assert_eq!(cache.len(), 6);
}
-/// A multi-token prefill followed by single-token steps — how a prompt is
-/// actually consumed — must also match.
+/// Prefill then single steps, the way a prompt is actually consumed.
#[test]
fn prefill_then_single_steps_matches_full_forward() {
let stack = stack();
@@ -59,7 +56,6 @@ fn prefill_then_single_steps_matches_full_forward() {
}
}
-/// Clearing the cache must give the same results as a fresh one.
#[test]
fn cleared_cache_starts_over() {
let stack = stack();
@@ -92,7 +88,7 @@ fn top_k_of_one_is_greedy_at_any_temperature() {
#[test]
fn tiny_top_p_is_greedy() {
- // The most likely token alone already exceeds p, so it is the whole nucleus.
+ // The top token alone already exceeds p.
let logits = Tensor::from_vec(vec![0.1, 3.0, -2.0, 1.5], &[4]);
let sampler = Sampler::new().top_p(0.01);
for _ in 0..20 {
@@ -114,7 +110,7 @@ fn sampling_is_reproducible_under_manual_seed() {
#[test]
fn sampling_respects_the_distribution() {
- // With one logit far above the rest, anything else is (2e-9)-unlikely.
+ // Anything else is about 2e-9 likely.
let logits = Tensor::from_vec(vec![10.0, -10.0, -10.0], &[3]);
let sampler = Sampler::new();
for _ in 0..100 {
@@ -122,8 +118,7 @@ fn sampling_respects_the_distribution() {
}
}
-/// A masked (padded) key position must have no influence on other positions:
-/// change its content and every unpadded row's output stays identical.
+/// Changing a padded position must not change the unpadded outputs.
#[test]
fn padding_mask_blocks_padded_positions() {
let encoder = TransformerStack::encoder(16, 2, 32, 1, 0.1);
@@ -131,7 +126,6 @@ fn padding_mask_blocks_padded_positions() {
let mask = Tensor::from_vec(vec![1.0, 1.0, 1.0, 0.0], &[1, 4]);
let a = Tensor::randn(&[1, 4, 16]);
- // Same first three positions, different content in the padded fourth.
let mut data = a.to_vec();
for slot in data.iter_mut().skip(3 * 16) {
*slot += 5.0;
@@ -147,8 +141,7 @@ fn padding_mask_blocks_padded_positions() {
);
}
-/// Same property through the decoder's cross-attention: masked memory
-/// positions must not reach the target side at all.
+/// Same through the decoder's cross-attention.
#[test]
fn memory_mask_blocks_padded_source_positions() {
let decoder = TransformerDecoder::new(16, 2, 32, 1, 0.1);
@@ -172,16 +165,15 @@ fn memory_mask_blocks_padded_source_positions() {
);
}
-/// The best second token can hide behind the second-best first token; a width
-/// of 2 must find the sequence a greedy width of 1 misses.
+/// The best path can start with the second-best token.
#[test]
fn beam_search_finds_what_greedy_misses() {
let step = |seq: &[usize]| match seq.last() {
- // From the start: token 0 is slightly likelier than token 1.
+ // Token 0 is slightly likelier than token 1 at the start.
Some(9) => vec![0.55f32.ln(), 0.45f32.ln(), f32::NEG_INFINITY],
- // After the greedy pick, everything is mediocre.
+ // After the greedy pick everything is mediocre.
Some(0) => vec![0.5f32.ln(), 0.5f32.ln(), f32::NEG_INFINITY],
- // After the runner-up, one continuation is near-certain.
+ // After the runner-up one continuation is near-certain.
_ => vec![0.99f32.ln(), 0.01f32.ln(), f32::NEG_INFINITY],
};
@@ -197,7 +189,7 @@ fn beam_search_finds_what_greedy_misses() {
#[test]
fn beam_search_stops_at_eos() {
- // Token 2 is end-of-sequence and always the most likely continuation.
+ // Token 2 is eos and always the most likely.
let step = |_: &[usize]| vec![0.0, 1.0, 5.0];
let out = BeamSearch::new(2).eos(2).decode(&[0], 10, step);
assert!(
diff --git a/tests/gradcheck.rs b/tests/gradcheck.rs
index 7376c2a..a016288 100644
--- a/tests/gradcheck.rs
+++ b/tests/gradcheck.rs
@@ -1,33 +1,18 @@
-//! Finite-difference gradient checks.
-//!
-//! Every backward rule is checked against central differences of its own
-//! forward. This is the safety net for the whole autograd layer: a rule with the
-//! wrong sign, a missing factor, or a forgotten broadcast reduction fails here
-//! rather than showing up later as a model that mysteriously will not converge.
-//!
-//! **Any new `Backward` rule needs a test in this file.**
-//!
-//! cargo test --no-default-features --test gradcheck
-
-// Test inputs are always passed as a slice. `&[x.clone()]` for a single input
-// costs one clone of a refcounted handle and keeps every call site reading the
-// same way as the multi-input ones.
+//! Every backward rule against central differences of its own forward.
+//! A new `Backward` rule needs a case here.
+
+// `&[x.clone()]` is a cheap refcount bump and keeps single and multi input calls uniform.
#![allow(clippy::cloned_ref_to_slice_refs)]
use fastnn::nn::Param;
use fastnn::prelude::*;
-/// Step size for the numeric derivative.
-///
-/// `f32` carries about seven digits and central differences spend roughly half
-/// of them. 1e-3 sits at the bottom of that trade-off: smaller drowns in
-/// rounding, larger starts measuring curvature instead of slope.
+/// Smaller drowns in f32 rounding, larger measures curvature instead of slope.
const STEP: f32 = 1e-3;
const TOLERANCE: f32 = 2e-2;
+const NORM_EPS: f32 = 1e-5;
-/// Check that `f`'s analytic gradients match its numeric ones at `values`.
-///
-/// `f` must produce a scalar — compose with `.sum()` if it does not.
+/// `f` must return a scalar.
#[track_caller]
fn gradcheck(name: &str, values: &[Tensor], f: impl Fn(&[Tensor]) -> Tensor) {
let params: Vec = values.iter().map(|v| Param::new(v.clone())).collect();
@@ -71,13 +56,9 @@ fn gradcheck(name: &str, values: &[Tensor], f: impl Fn(&[Tensor]) -> Tensor) {
}
}
-/// Deterministic, well-spread values in roughly `(-0.9, 0.9)`.
-///
-/// Bounded on purpose. Central differences subtract two nearby loss values, so
-/// if the loss can grow into the thousands the difference loses most of `f32`'s
-/// digits to cancellation and the check fails on its own arithmetic. The golden
-/// ratio spreads the values without repeats, and none land on 0 or ±1, where a
-/// wrong sign or exponent could still give the right answer.
+/// Values kept small on purpose: a loss in the thousands loses most of f32's
+/// digits to cancellation in the difference. The golden ratio avoids repeats
+/// and keeps values off 0 and ±1, where a wrong sign or exponent can still pass.
fn spread(n: usize, low: f32, high: f32) -> Vec {
const GOLDEN: f32 = 0.618_034;
(0..n)
@@ -89,13 +70,11 @@ fn sample(shape: &[usize]) -> Tensor {
Tensor::from_vec(spread(shape.iter().product(), -0.9, 0.9), shape)
}
-/// Strictly positive values, for ops undefined at or below zero.
+/// For ops undefined at or below zero.
fn positive(shape: &[usize]) -> Tensor {
Tensor::from_vec(spread(shape.iter().product(), 0.3, 1.4), shape)
}
-// ── Arithmetic ───────────────────────────────────────────────────────────────
-
#[test]
fn arithmetic() {
let a = sample(&[2, 3]);
@@ -118,8 +97,7 @@ fn arithmetic() {
#[test]
fn broadcasting_reduces_gradients() {
- // The [1, 3] operand is stretched over 4 rows, so its gradient must be the
- // sum of all four contributions — the classic bias-gradient case.
+ // The [1, 3] operand is stretched over 4 rows, so its gradient sums four terms.
gradcheck("broadcast add", &[sample(&[4, 3]), sample(&[1, 3])], |x| {
x[0].add(&x[1]).sum()
});
@@ -131,8 +109,6 @@ fn broadcasting_reduces_gradients() {
});
}
-// ── Element-wise maths ───────────────────────────────────────────────────────
-
#[test]
fn unary_maths() {
let x = sample(&[2, 3]);
@@ -144,13 +120,10 @@ fn unary_maths() {
gradcheck("sqrt", &[positive(&[2, 3])], |x| x[0].sqrt().sum());
gradcheck("powf", &[positive(&[2, 3])], |x| x[0].powf(2.5).sum());
gradcheck("abs", &[x.clone()], |x| x[0].abs().sum());
- // Bounds sit clear of every sample value: at a clamp boundary the true
- // derivative jumps, and central differences would straddle both sides.
+ // Bounds stay clear of the samples, the derivative jumps at a clamp edge.
gradcheck("clamp", &[x], |x| x[0].clamp(-9.0, 9.0).sum());
}
-// ── Activations ──────────────────────────────────────────────────────────────
-
#[test]
fn activations() {
let x = sample(&[3, 4]);
@@ -168,8 +141,7 @@ fn softmax_family() {
let x = sample(&[3, 4]);
let weights = sample(&[3, 4]);
- // Weighted, not a plain sum: softmax rows total 1 whatever the input, so an
- // unweighted sum has zero gradient and would pass against any rule at all.
+ // Softmax rows sum to 1, so a plain sum has zero gradient and passes any rule.
let w = weights.clone();
gradcheck("softmax", &[x.clone()], move |x| {
x[0].softmax().mul(&w).sum()
@@ -180,8 +152,6 @@ fn softmax_family() {
});
}
-// ── Matrix multiplication ────────────────────────────────────────────────────
-
#[test]
fn matmul_layouts() {
gradcheck("matmul", &[sample(&[2, 3]), sample(&[3, 4])], |x| {
@@ -208,8 +178,6 @@ fn matmul_batched() {
);
}
-// ── Reductions ───────────────────────────────────────────────────────────────
-
#[test]
fn reductions() {
let x = sample(&[3, 4]);
@@ -228,8 +196,6 @@ fn reductions() {
});
}
-// ── Views ────────────────────────────────────────────────────────────────────
-
#[test]
fn views() {
let x = sample(&[2, 3, 4]);
@@ -239,8 +205,7 @@ fn views() {
x[0].reshape(&[-1]).mul(&w).sum()
});
- // Weighted so the reordering matters: a plain sum is invariant to any
- // permutation and would pass even if the inverse were computed wrongly.
+ // A plain sum is blind to permutations.
let w = sample(&[4, 3, 2]);
gradcheck("permute", &[x.clone()], move |x| {
x[0].permute(&[2, 1, 0]).mul(&w).sum()
@@ -263,7 +228,7 @@ fn views() {
#[test]
fn narrow_and_join() {
- // Only part of the input reaches the loss; the rest must get exactly zero.
+ // The rest of the input must get exactly zero.
let w = sample(&[2, 3]);
gradcheck("narrow", &[sample(&[2, 6])], move |x| {
x[0].narrow(1, 2, 3).mul(&w).sum()
@@ -280,29 +245,23 @@ fn narrow_and_join() {
});
}
-// ── Indexing ─────────────────────────────────────────────────────────────────
-
#[test]
fn index_select_accumulates_repeats() {
- // Row 1 is picked twice, so its gradient must be the sum of both lookups.
+ // Row 1 is picked twice.
let w = sample(&[4, 3]);
gradcheck("index_select", &[sample(&[5, 3])], move |x| {
x[0].index_select(&[0, 1, 1, 4]).mul(&w).sum()
});
}
-// ── Structured ops ───────────────────────────────────────────────────────────
-
#[test]
fn convolution() {
- // [1, 1, 4, 4] with a 2x2 stride-2 window -> [1, 1*2*2 patch, 2*2 positions].
let w = sample(&[1, 4, 4]);
gradcheck("im2col", &[sample(&[1, 1, 4, 4])], move |x| {
x[0].im2col(Window::square(2, 2, 0)).mul(&w).sum()
});
- // The whole layer end to end: im2col and matmul must compose correctly, and
- // overlapping windows must accumulate into the same input pixels.
+ // Overlapping windows must accumulate into the same input pixels.
gradcheck(
"conv2d",
&[sample(&[2, 2, 4, 4]), sample(&[3, 2, 3, 3]), sample(&[3])],
@@ -326,7 +285,7 @@ fn pooling() {
x[0].avg_pool2d(Window::square(2, 2, 0)).mul(&w).sum()
});
- // Uneven spans: 5 does not divide by 2, so the windows differ in size.
+ // 5 does not divide by 2, so the windows differ in size.
let w = sample(&[1, 2, 2, 2]);
gradcheck("adaptive_avg_pool2d", &[sample(&[1, 2, 5, 5])], move |x| {
x[0].adaptive_avg_pool2d((2, 2)).mul(&w).sum()
@@ -339,7 +298,7 @@ fn normalization() {
gradcheck(
"layer_norm",
&[sample(&[3, 4]), Tensor::ones(&[4]), Tensor::zeros(&[4])],
- move |x| x[0].layer_norm(&x[1], &x[2], 1e-5).mul(&w).sum(),
+ move |x| x[0].layer_norm(&x[1], &x[2], NORM_EPS).mul(&w).sum(),
);
let w = sample(&[2, 2, 2, 2]);
@@ -350,7 +309,7 @@ fn normalization() {
Tensor::ones(&[2]),
Tensor::zeros(&[2]),
],
- move |x| x[0].batch_norm2d(&x[1], &x[2], 1e-5).0.mul(&w).sum(),
+ move |x| x[0].batch_norm2d(&x[1], &x[2], NORM_EPS).0.mul(&w).sum(),
);
let w = sample(&[3, 4]);
@@ -360,8 +319,6 @@ fn normalization() {
});
}
-// ── Losses ───────────────────────────────────────────────────────────────────
-
#[test]
fn losses() {
gradcheck("cross_entropy", &[sample(&[3, 4])], |x| {
@@ -387,8 +344,6 @@ fn losses() {
});
}
-// ── Whole layers ─────────────────────────────────────────────────────────────
-
#[test]
fn a_full_model_trains() {
manual_seed(7);
@@ -465,8 +420,7 @@ fn gradients_accumulate_until_zeroed() {
#[test]
fn a_parameter_used_twice_sums_both_paths() {
- // Weight tying: the same handle appears in two places, so its gradient must
- // be the sum of both. This works because both uses share one GradSlot.
+ // Weight tying: both uses share one GradSlot.
let shared = Param::new(sample(&[2, 2]));
let x = sample(&[2, 2]);
@@ -527,8 +481,7 @@ fn masked_attention_matches_finite_differences() {
});
}
-/// The configurable cross-entropy has its own fused backward; check it under
-/// every option at once — smoothing, class weights, and an ignored row.
+/// Fused backward, with every option on at once.
#[test]
fn weighted_cross_entropy_matches_finite_differences() {
let loss_fn = CrossEntropyLoss::new()
@@ -541,8 +494,7 @@ fn weighted_cross_entropy_matches_finite_differences() {
});
}
-/// The per-row reduction takes a `[batch]` upstream gradient — the other
-/// branch of the backward rule, checked with a weighted combination.
+/// Takes a `[batch]` upstream gradient, the other branch of the rule.
#[test]
fn per_row_cross_entropy_matches_finite_differences() {
let weight = Tensor::from_vec(vec![0.7, -0.3, 1.2], &[3]);
@@ -553,13 +505,9 @@ fn per_row_cross_entropy_matches_finite_differences() {
});
}
-/// col2im is a new forward op; check it (and its Im2ColBackward) directly,
-/// weighted so overlap accumulation cannot hide behind a symmetric sum.
#[test]
fn col2im_matches_finite_differences() {
- use fastnn::tensor::Window;
let window = Window::square(2, 2, 0);
- // Columns for a 4×4 target image: patch = 1·2·2, positions = 2·2.
let weight = sample(&[1, 1, 4, 4]);
gradcheck("col2im", &[sample(&[1, 4, 4])], |v| {
@@ -567,12 +515,9 @@ fn col2im_matches_finite_differences() {
});
}
-/// Dilation changes the index map; the finite-difference check proves the
-/// backward folds along the same stretched footprint the forward reads.
#[test]
fn dilated_im2col_matches_finite_differences() {
- use fastnn::tensor::Window;
- // Span (2−1)·2+1 = 3 on a padded 5×5 input → 5×5 = 25 output positions.
+ // Span 3 on a padded 5x5 input gives 25 output positions.
let window = Window::square(2, 1, 1).dilated(2);
let weight = sample(&[1, 4, 25]);
diff --git a/tests/losses.rs b/tests/losses.rs
index 26b7d6f..e9e3459 100644
--- a/tests/losses.rs
+++ b/tests/losses.rs
@@ -1,9 +1,4 @@
-//! Loss functions against hand-computed references.
-//!
-//! Each case pins a loss to a value worked out on paper from its defining
-//! formula, so a refactor that changes semantics — a swapped sign, a missing
-//! normalizer, a wrong direction of KL — fails loudly with numbers a reader
-//! can recompute.
+//! Losses against values worked out by hand from their formulas.
use fastnn::nn::{
contrastive, cosine_embedding, dice, focal_bce_with_logits, gaussian_nll, info_nce,
@@ -11,21 +6,26 @@ use fastnn::nn::{
};
use fastnn::prelude::*;
+const TOLERANCE: f32 = 1e-4;
+
#[track_caller]
fn assert_close(got: f32, want: f32, what: &str) {
- assert!((got - want).abs() < 1e-4, "{what}: got {got}, want {want}");
+ assert!(
+ (got - want).abs() < TOLERANCE,
+ "{what}: got {got}, want {want}"
+ );
}
#[test]
fn huber_is_quadratic_inside_delta_and_linear_outside() {
let prediction = Tensor::from_vec(vec![0.5, 3.0], &[2, 1]);
let target = Tensor::zeros(&[2, 1]);
- // ½·0.5² = 0.125 inside; δ(|3| − ½δ) = 2.5 outside; mean = 1.3125.
+ // 0.5 * 0.5^2 = 0.125 inside, 1 * (3 - 0.5) = 2.5 outside, mean 1.3125.
assert_close(huber(&prediction, &target, 1.0).item(), 1.3125, "huber");
assert_close(
smooth_l1(&prediction, &target).item(),
1.3125,
- "smooth_l1 = huber(δ=1)",
+ "smooth_l1 = huber(delta=1)",
);
}
@@ -33,9 +33,8 @@ fn huber_is_quadratic_inside_delta_and_linear_outside() {
fn kl_divergence_matches_the_definition() {
let target = Tensor::from_vec(vec![0.5, 0.5], &[1, 2]);
let prediction = Tensor::from_vec(vec![0.25, 0.75], &[1, 2]);
- // 0.5·ln(0.5/0.25) + 0.5·ln(0.5/0.75) = 0.5(ln 2 + ln ⅔) ≈ 0.14384.
+ // 0.5 * ln(0.5 / 0.25) + 0.5 * ln(0.5 / 0.75) = 0.14384
assert_close(kl_divergence(&prediction, &target).item(), 0.14384, "kl");
- // KL(p ‖ p) = 0.
assert_close(
kl_divergence(&target, &target).item(),
0.0,
@@ -56,12 +55,12 @@ fn nll_of_log_softmax_equals_fused_cross_entropy() {
#[test]
fn gaussian_and_poisson_nll_match_their_formulas() {
- // ½(log σ² + (x−μ)²/σ²) with μ=0, log σ²=0, x=2 → ½·4 = 2.
+ // 0.5 * (logvar + (x - mean)^2 / var) = 0.5 * 4
let (mean, logvar) = (Tensor::zeros(&[1]), Tensor::zeros(&[1]));
let x = Tensor::from_vec(vec![2.0], &[1]);
assert_close(gaussian_nll(&mean, &logvar, &x).item(), 2.0, "gaussian nll");
- // e^{log λ} − k·log λ with log λ = 0, k = 3 → 1.
+ // exp(log_rate) - count * log_rate = 1
let log_rate = Tensor::zeros(&[1]);
let count = Tensor::from_vec(vec![3.0], &[1]);
assert_close(poisson_nll(&log_rate, &count).item(), 1.0, "poisson nll");
@@ -71,11 +70,11 @@ fn gaussian_and_poisson_nll_match_their_formulas() {
fn focal_with_gamma_zero_is_alpha_scaled_bce() {
let logits = Tensor::from_vec(vec![0.7, -1.3, 0.2], &[3, 1]);
let target = Tensor::from_vec(vec![1.0, 0.0, 1.0], &[3, 1]);
- // γ = 0 removes the modulation; α = ½ scales both classes equally.
+ // gamma 0 drops the modulation, alpha 0.5 scales both classes equally.
assert_close(
focal_bce_with_logits(&logits, &target, 0.0, 0.5).item(),
0.5 * bce_with_logits(&logits, &target).item(),
- "focal(γ=0, α=½) = ½·bce",
+ "focal(gamma=0, alpha=0.5) = 0.5 * bce",
);
}
@@ -85,7 +84,7 @@ fn dice_is_zero_on_perfect_binary_overlap_and_high_on_none() {
assert_close(dice(&mask, &mask).item(), 0.0, "perfect overlap");
let inverse = Tensor::from_vec(vec![0.0, 1.0, 0.0, 0.0], &[4]);
- // No overlap: 1 − (0+1)/(4+1) = 0.8.
+ // 1 - (0 + 1) / (4 + 1)
assert_close(dice(&mask, &inverse).item(), 0.8, "no overlap");
}
@@ -102,8 +101,8 @@ fn plain_cross_entropy_loss_matches_the_fused_function() {
#[test]
fn label_smoothing_matches_the_smoothed_target_distribution() {
- // p = [¼, ¾], target 1, s = 0.2 ⇒ q = [0.1, 0.9]:
- // loss = −(0.1·ln ¼ + 0.9·ln ¾) ≈ 0.39752.
+ // p = [1/4, 3/4], target 1, smoothing 0.2 gives q = [0.1, 0.9]
+ // loss = -(0.1 * ln(1/4) + 0.9 * ln(3/4)) = 0.39752
let logits = Tensor::from_vec(vec![0.0, 3.0f32.ln()], &[1, 2]);
let loss = CrossEntropyLoss::new()
.label_smoothing(0.2)
@@ -122,11 +121,7 @@ fn class_weights_take_a_weighted_mean() {
.class_weights(vec![1.0, 3.0])
.compute(&logits, &[0, 1]);
let expected = (1.0 * separate[0] + 3.0 * separate[1]) / 4.0;
- assert_close(
- weighted.item(),
- expected,
- "weighted mean over Σw, not batch",
- );
+ assert_close(weighted.item(), expected, "mean over the weight sum");
}
#[test]
@@ -135,7 +130,6 @@ fn ignored_rows_contribute_neither_loss_nor_gradient() {
let loss_fn = CrossEntropyLoss::new().ignore_index(1);
let loss = loss_fn.compute(&logits.tensor(), &[0, 1]);
- // Only row 0 participates, so the mean is exactly row 0's loss.
assert_close(
loss.item(),
cross_entropy(&logits.value().narrow(0, 0, 1), &[0]).item(),
@@ -150,7 +144,7 @@ fn ignored_rows_contribute_neither_loss_nor_gradient() {
);
assert!(grad[0] != 0.0, "participating row must receive gradient");
- // Every row ignored: zero loss, and finite everywhere.
+ // Every row ignored.
let all_ignored = loss_fn.compute(&logits.tensor().detach(), &[1, 1]);
assert_close(all_ignored.item(), 0.0, "all-ignored batch");
}
@@ -170,13 +164,13 @@ fn reduction_modes_relate_as_mean_sum_and_rows() {
.compute(&logits, &targets);
assert_eq!(rows.shape(), &[2], "None keeps one loss per row");
- assert_close(sum, mean * 2.0, "sum = mean·batch");
+ assert_close(sum, mean * 2.0, "sum = mean * batch");
assert_close(rows.to_vec().iter().sum::(), sum, "rows sum to Sum");
}
#[test]
fn metric_losses_match_hand_worked_geometry() {
- // Triplet: d(a,p) = 1, d(a,n) = 1.5, margin 1 → hinge at 0.5.
+ // d(a,p) = 1, d(a,n) = 1.5, margin 1, so the hinge is 0.5.
let anchor = Tensor::from_vec(vec![0.0, 0.0], &[1, 2]);
let positive = Tensor::from_vec(vec![1.0, 0.0], &[1, 2]);
let negative = Tensor::from_vec(vec![1.5, 0.0], &[1, 2]);
@@ -186,8 +180,7 @@ fn metric_losses_match_hand_worked_geometry() {
"triplet",
);
- // Cosine: orthogonal vectors. Similar pair costs 1 − 0 = 1; dissimilar
- // pair with margin 0 costs relu(0 − 0) = 0; mean = 0.5.
+ // Orthogonal vectors: the similar pair costs 1, the dissimilar one 0.
let a = Tensor::from_vec(vec![1.0, 0.0, 1.0, 0.0], &[2, 2]);
let b = Tensor::from_vec(vec![0.0, 1.0, 0.0, 1.0], &[2, 2]);
assert_close(
@@ -196,7 +189,7 @@ fn metric_losses_match_hand_worked_geometry() {
"cosine embedding",
);
- // Ranking: x1 already margin ahead → 0; wrong order → 1 + margin.
+ // x1 already ahead costs 0, the wrong order costs 1 + margin.
let x1 = Tensor::from_vec(vec![2.0, 2.0], &[2, 1]);
let x2 = Tensor::from_vec(vec![1.0, 1.0], &[2, 1]);
assert_close(
@@ -205,8 +198,7 @@ fn metric_losses_match_hand_worked_geometry() {
"margin ranking",
);
- // Contrastive: matched at distance 0 → 0; mismatched at distance 1 with
- // margin 2 → (2−1)² = 1; mean = 0.5.
+ // Matched at distance 0 costs 0, mismatched at distance 1 with margin 2 costs 1.
let u = Tensor::from_vec(vec![0.0, 0.0, 0.0, 0.0], &[2, 2]);
let v = Tensor::from_vec(vec![0.0, 0.0, 1.0, 0.0], &[2, 2]);
assert_close(
@@ -218,8 +210,8 @@ fn metric_losses_match_hand_worked_geometry() {
#[test]
fn info_nce_on_an_identity_batch_matches_the_closed_form() {
- // Orthonormal queries equal to keys at τ = 1: each row's logits are its
- // own row of I, so loss = −ln(e/(e + (n−1))) = ln(1 + (n−1)e⁻¹).
+ // Queries equal keys at temperature 1, so each row's logits are a row of I
+ // and the loss is ln(1 + (n - 1) / e).
let identity = Tensor::from_vec(vec![1.0, 0.0, 0.0, 1.0], &[2, 2]);
let expected = (1.0f32 + 1.0 * (-1.0f32).exp()).ln();
assert_close(
diff --git a/tests/robustness.rs b/tests/robustness.rs
index c8fb778..dec40e9 100644
--- a/tests/robustness.rs
+++ b/tests/robustness.rs
@@ -1,16 +1,10 @@
-//! The guard rails a long unattended run depends on.
-//!
-//! Each of these covers a failure that is silent without it: an optimizer that
-//! forgets its momentum, a NaN that only surfaces once the checkpoint is already
-//! poisoned, a frozen layer that trains anyway.
-//!
-//! cargo test --no-default-features --test robustness
+//! Resuming, anomaly detection, freezing and the fallible ops.
use fastnn::autograd::detect_anomaly;
use fastnn::prelude::*;
use fastnn::serialize::{load_training, save_training};
-/// A scratch path that cleans itself up.
+/// Deleted on drop.
struct TempFile(std::path::PathBuf);
impl TempFile {
@@ -56,8 +50,6 @@ fn train_steps(
}
}
-// ── Float indices ────────────────────────────────────────────────────────────
-
fn assert_invalid_float_indices(mut consume: impl FnMut(f32)) {
let upper = (usize::MAX as f64 + 1.0) as f32;
let mut failures = Vec::new();
@@ -146,21 +138,19 @@ fn float_indices_labels_accept_largest_fitting_integer() {
assert_eq!(batch.labels(), vec![largest as usize]);
}
-// ── Resuming a run ───────────────────────────────────────────────────────────
-
#[test]
fn resuming_matches_training_straight_through() {
let file = TempFile::new("resume");
let (inputs, targets) = batch();
- // Run twenty steps without interruption.
+ // Twenty steps straight through.
manual_seed(9);
let reference = small_model();
let mut opt = Adam::new(reference.parameters(), 0.05);
train_steps(&reference, &mut opt, &inputs, &targets, 20);
let uninterrupted = cross_entropy(&reference.forward(&inputs), &targets).item();
- // The same twenty steps, saved and reloaded halfway.
+ // The same twenty with a save and reload at ten.
manual_seed(9);
let resumed = small_model();
let mut opt = Adam::new(resumed.parameters(), 0.05);
@@ -182,10 +172,8 @@ fn resuming_matches_training_straight_through() {
#[test]
fn a_cold_optimizer_takes_a_different_path() {
- // Guards the test above. If the moments and step count did not really survive
- // the round trip, both paths here would be identical and
- // `resuming_matches_training_straight_through` could pass while restoring
- // nothing. Restoring the weights alone must visibly differ.
+ // Guards the test above: if no optimizer state were restored, the two paths
+ // would match and that test could pass while restoring nothing.
let file = TempFile::new("cold");
let (inputs, targets) = batch();
@@ -198,7 +186,7 @@ fn a_cold_optimizer_takes_a_different_path() {
let continue_step = |restore_optimizer: bool| {
let restored = small_model();
let mut opt = Adam::new(restored.parameters(), 0.05);
- // Both paths load the same weights; only one also loads the optimizer.
+ // Both load the same weights, only one keeps the loaded optimizer.
let mut loader = Adam::new(restored.parameters(), 0.05);
load_training(&restored, &mut loader, file.path()).unwrap();
if restore_optimizer {
@@ -264,14 +252,12 @@ fn a_mismatched_optimizer_is_reported() {
train_steps(&model, &mut opt, &inputs, &targets, 3);
save_training(&model, &opt, 3, file.path()).unwrap();
- // Same architecture, but the optimizer covers only part of it.
+ // The optimizer covers only part of the model.
let other = small_model();
let mut fewer = Adam::new(other.parameters()[..2].to_vec(), 0.05);
assert!(load_training(&other, &mut fewer, file.path()).is_err());
}
-// ── Anomaly detection ────────────────────────────────────────────────────────
-
#[test]
#[should_panic(expected = "Log")]
fn anomaly_detection_names_the_op_at_fault() {
@@ -313,8 +299,6 @@ fn finiteness_guard_spots_a_poisoned_gradient() {
assert!(!opt.gradients_are_finite());
}
-// ── Freezing ─────────────────────────────────────────────────────────────────
-
#[test]
fn a_frozen_layer_does_not_move() {
let backbone = Linear::new(4, 6);
@@ -351,8 +335,7 @@ fn a_frozen_layer_does_not_move() {
#[test]
fn freezing_skips_the_gradient_entirely() {
- // A frozen parameter hands out a detached tensor, so the backward pass stops
- // there — no gradient is computed only to be thrown away.
+ // A frozen parameter hands out a detached tensor, so no gradient is computed.
let layer = Linear::new(4, 6);
layer.freeze();
@@ -388,8 +371,6 @@ fn freezing_reaches_through_a_container() {
);
}
-// ── Fallible ops ─────────────────────────────────────────────────────────────
-
#[test]
fn shape_errors_are_reported_not_fatal() {
let x = Tensor::zeros(&[4, 8]);
diff --git a/tests/safetensors.rs b/tests/safetensors.rs
index 496c4ff..89de101 100644
--- a/tests/safetensors.rs
+++ b/tests/safetensors.rs
@@ -1,4 +1,4 @@
-//! The safetensors reader and writer, end to end.
+//! Safetensors round trips, renaming, half precision and corrupt files.
use std::collections::HashMap;
@@ -47,7 +47,7 @@ fn partial_load_reports_instead_of_guessing() {
let path = temp("fastnn_st_partial.safetensors");
save_safetensors(&small, &path).unwrap();
- // A model with an extra layer: the first layer matches, the second cannot.
+ // The extra layer has nothing to load.
let bigger = Sequential::new()
.add(Linear::new(4, 4))
.add(Linear::new(4, 3));
@@ -66,7 +66,7 @@ fn renaming_maps_foreign_keys_onto_the_model() {
let path = temp("fastnn_st_rename.safetensors");
save_safetensors(&source, &path).unwrap();
- // A container gives the same layer different names; the map bridges them.
+ // The container names the layer differently.
let target = Sequential::new().add(Linear::new(3, 2));
let rename: HashMap = [
("weight".to_string(), "0.weight".to_string()),
@@ -105,8 +105,7 @@ fn shape_clash_is_an_error_that_names_the_tensor() {
);
}
-/// A file written by hand in the two half-precision dtypes: loading must widen
-/// each stored value to exactly the f32 the conversion tables promise.
+/// A hand-written file in both half dtypes, every value must widen exactly.
#[test]
fn half_precision_tensors_widen_exactly() {
let f16_values = [1.0f32, -2.5, 0.15625];
@@ -163,7 +162,7 @@ fn corrupt_headers_fail_loudly_not_quietly() {
let header = "{\"w\":{\"dtype\":\"F32\",\"shape\":[4],\"data_offsets\":[0,16]}}";
let mut f = (header.len() as u64).to_le_bytes().to_vec();
f.extend_from_slice(header.as_bytes());
- f.extend_from_slice(&[0u8; 4]); // only 4 of the claimed 16 bytes
+ f.extend_from_slice(&[0u8; 4]); // header claims 16
f
}),
("shape and span disagree", {