Numerics, statistics and neural-network primitives for zimr, with a CPU
implementation and a matching set of GPU kernels. This document covers the library as it exists in
src/zimrnum.zig. Sections appear when the code they describe does.
A numerical library for zimr: tensors and the operations over them, statistics, and the
primitives a neural network is built from. Every operation has a CPU implementation in
src/zimrnum.zig, and most have a matching GPU kernel checked against it.
src/zimrnum.zig imports std, zm
(src/zimrmath.zig) and kompute (src/kompute.zig), and
nothing else. Its whole closure is free of shaders, of the WebGPU bindings and of the engine, so
it compiles and tests in seconds, and a module that is itself zm-only can adopt it
without acquiring a GPU dependency. The constraint is enforced by build.zig, which
creates the zn module with exactly those two imports.
The CPU implementation is the reference. Each GPU kernel is authored in Zig against
kompute, compiled to SPIR-V by the stock Zig compiler and translated to WGSL by
src/spv2wgsl.zig; every kernel is compared element by element against the
corresponding zn function on a real device. §13 covers what that
comparison establishes and what it cannot.
Where both sides need the same mathematical function, both call zm. That matters
because zm selects its own implementation per backend — zm.atan2 is
std.math.atan2 on the host and a polynomial on the GPU — so calling it keeps the
choice in one place.
| Area | Present |
|---|---|
Base: Error, float guards, Ddof, approxEqAbs | yes |
Random: counter-based Rng, uniform, normal, bounded integers | yes |
Memory: Ctx, arenas | yes |
| Tensors: views, slicing, permutation, broadcasting, copies | yes |
| Elementwise: arithmetic, masks, rounding, transcendentals, activations, gradients | yes |
| Reductions: sum, mean, min, max, along an axis | yes |
Linear algebra: matmul | yes |
| Network pieces: softmax, layer norm, MSE, SGD step | yes |
| Statistics: variance, deviation, covariance, correlation, scaling | yes |
| Autograd | no |
| Layers, optimisers beyond SGD, a training loop | no |
| Decompositions: LU, Cholesky, QR, eigh, SVD | no |
| FFT, convolution, pooling | no |
| Reinforcement learning, dataframes, serialisation | no |
| Element types other than f16/f32/f64; sparse tensors | no |
This fits a line to noisy data and reports the loss before and after. It is the whole library in miniature: a tensor, a layer, a tape, a gradient, an update.
const std = @import("std");
const zn = @import("zn");
// A complete program: fit a line to noisy data, then report the error before and after.
pub fn main() !void {
var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator);
defer arena.deinit();
const gpa: std.mem.Allocator = arena.allocator();
// 1. Some data. y = 3x + 2, with noise.
const rng: zn.Rng = .init(20260913);
const x: zn.Tensor(f64) = try zn.linspace(f64, gpa, 0, 1, 32, .include);
const y: zn.Tensor(f64) = try zn.Tensor(f64).alloc(gpa, &.{ 32, 1 });
const noise: zn.Tensor(f64) = try zn.Tensor(f64).alloc(gpa, &.{32});
rng.fillNormal(f64, noise.data);
for (0..32) |i| {
y.data[i] = 3.0 * x.data[i] + 2.0 + 0.1 * noise.data[i];
}
const inputs: zn.Tensor(f64) = try x.reshape(&.{ 32, 1 });
// 2. A layer to fit with.
const line = try zn.Dense(f64).init(gpa, rng.split(1), 1, 1, .xavier);
// 3. Train.
var first: f64 = 0;
var last: f64 = 0;
for (0..400) |step| {
var graph: zn.Graph(f64) = .init(gpa);
const guess: zn.Dense(f64).Attached = try line.attach(
&graph,
try graph.constant(inputs),
);
const loss: zn.Var = try graph.mseLoss(guess.out, try graph.constant(y));
const value: f64 = graph.valueOf(loss).data[0];
if (step == 0) {
first = value;
}
last = value;
try graph.backward(loss);
const dw: zn.Tensor(f64) = try graph.gradOf(guess.weight);
const db: zn.Tensor(f64) = try graph.gradOf(guess.bias);
for (line.weight.data, dw.data) |*w, g| {
w.* -= 0.1 * g;
}
for (line.bias.data, db.data) |*b, g| {
b.* -= 0.1 * g;
}
}
std.debug.print("\n loss {d:.4} -> {d:.4}\n", .{ first, last });
std.debug.print(
" slope {d:.3} (true 3.0), intercept {d:.3} (true 2.0)\n",
.{ line.weight.data[0], line.bias.data[0] },
);
}Rng/// A reproducible random source addressed by (seed, index) rather than by mutation.
///
/// WHY COUNTER-BASED, AND NOT A STATEFUL GENERATOR
///
/// A classic PRNG carries state and advances it, so drawing the i-th value means drawing the
/// first i. That is fine on one CPU thread and useless everywhere else this library is going:
///
/// * **A GPU kernel has no shared state.** Thread `i` must be able to compute ITS value without
/// any other thread having run. A counter-based generator is a pure function of the index, so
/// the kernel is `value(seed, thread_id)` and needs no sequencing at all.
/// * **CPU and GPU must agree EXACTLY.** The CPU path is this library's oracle (plan section 5 rung
/// 1); an oracle that produces a different random stream than the thing it checks is not an
/// oracle. Same function, same integers, same answer.
/// * **A failing case is reproducible from two numbers.** "seed 7, index 4096" is the whole
/// state, so a bad draw can be re-created in a unit test without replaying a sequence.
///
/// THE MIXER IS DELIBERATELY MODEST AND THE TESTS SAY SO. `mix32` is the murmur3 finalizer -
/// a bijection with good avalanche, applied to `index` folded with `seed`. It is NOT a
/// cryptographic generator and NOT claimed to pass a full statistical battery; what it is pinned
/// on is the properties this library actually relies on: determinism, range, stream independence,
/// and first/second moments. A stronger generator can replace the body without touching a single
/// call site, precisely because the interface is (seed, index) rather than a stream.
pub const Rng = struct {
seed: u32,
/// A generator for a given seed. Two `Rng`s with the same seed are the same function.
pub fn init(seed: u32) Rng {
return .{ .seed = seed };
}
/// A DERIVED generator, so one seed can feed several independent uses without the caller
/// inventing seeds. `parent.split(0)` and `parent.split(1)` are unrelated streams.
///
/// Mixing the label rather than adding it: `seed + 1` and `seed + 2` are adjacent inputs to
/// the mixer, and adjacency is exactly what a hash is asked not to preserve. Passing the
/// label through `mix32` first means the derived seeds are as far apart as any other pair.
pub fn split(self: Rng, label: u32) Rng {
return .{ .seed = mix32(self.seed ^ mix32(label ^ 0x9e3779b9)) };
}
/// The number of mantissa bits `T` can hold, including the implicit leading one. This is how
/// many random bits a draw in `[0, 1)` may use: fewer wastes range, MORE ROUNDS UP TO 1.0 and
/// breaks the half-open interval.
fn mantissaWidth(comptime T: type) u32 {
comptime requireFloat(T);
return switch (T) {
f16 => 11,
f32 => 24,
f64 => 53,
else => unreachable,
};
}
/// One 32-bit word of the generator's output at `index`, from independent stream `stream`.
///
/// EACH KIND OF DRAW OWNS A STREAM, SO KINDS DO NOT INTERFERE
///
/// Without this, every method would start from `bits(index)` and drawing two different KINDS
/// at one index would return values that determine each other - `intBelow(i, n)` would be a
/// function of `unitFloat(i)`. A caller has no reason to expect that and every reason to be
/// hurt by it, so the stream is part of the addressing rather than the caller's problem.
///
/// The stream index is folded into the MIXER, not into the index. Folding it into the index
/// (`bits(index ^ K)`) makes indices `i` and `i ^ K` draw from the same two words with their
/// roles swapped, so those pairs are not independent - which is exactly the defect this
/// replaced.
///
/// `stream == 0` adds nothing, so `bits` is unchanged by the introduction of rounds.
fn word(self: Rng, index: u32, stream: u32) u32 {
return mix32(mix32(index *% 0x9e3779b9) ^ self.seed ^ (stream *% 0x85ebca6b));
}
/// The raw 32-bit draw at `index`, stream 0. Distinct from every other kind of draw.
pub fn bits(self: Rng, index: u32) u32 {
return self.word(index, 0);
}
/// A float in `[0, 1)` from stream `stream` (and `stream + 1` when `T` needs more than 32 bits).
///
/// Exactly `mantissaWidth(T)` random bits, scaled by 2^-width. The largest result is
/// (2^width - 1) / 2^width, which is the largest `T` strictly below one - so the interval is
/// half-open for every width, including `f16`, where using 24 bits rounds to 1.0 instead.
fn unitFloatRound(self: Rng, comptime T: type, index: u32, comptime stream: u32) T {
comptime requireFloat(T);
const width: u32 = comptime mantissaWidth(T);
const first: u64 = self.word(index, stream);
const mantissa: u64 = if (comptime width <= 32)
first >> comptime @as(u6, @intCast(32 - width))
else
(first << comptime @as(u6, @intCast(width - 32))) |
(@as(u64, self.word(index, stream + 1)) >> comptime @as(u6, @intCast(64 - width)));
const numerator: T = @floatFromInt(mantissa);
const denominator: T = @floatFromInt(@as(u64, 1) << comptime @as(u6, @intCast(width)));
return numerator / denominator;
}
/// A float in `[0, 1)`, drawn at `index`.
pub fn unitFloat(self: Rng, comptime T: type, index: u32) T {
return self.unitFloatRound(T, index, 1);
}
/// A float in `[low, high)`, drawn at `index`.
pub fn uniform(
self: Rng,
comptime T: type,
index: u32,
low: T,
high: T,
) T {
comptime requireFloat(T);
return low + (high - low) * self.unitFloat(T, index);
}
/// A standard normal draw (mean 0, variance 1) at `index`, by the Box-Muller transform.
///
/// ITS TWO UNIFORMS COME FROM ITS OWN STREAMS (3-6), not from the uniform stream, so
/// `normal(i)` and `unitFloat(i)` are independent and a caller mixing them needs no care.
///
/// Box-Muller rather than a ziggurat because a ziggurat REJECTS: the number of uniforms it
/// consumes varies per sample, which would make the value at `index` depend on how many
/// rejections happened. That is not a pure function of the index and would not survive being
/// evaluated on an independent GPU thread. Box-Muller consumes exactly two, always.
///
/// The radius draw is taken from `(0, 1]` rather than `[0, 1)` because `@log(0)` is -inf.
pub fn normal(self: Rng, comptime T: type, index: u32) T {
comptime requireFloat(T);
const two_pi: T = 6.283185307179586;
const radius_draw: T = 1.0 - self.unitFloatRound(T, index, 3);
const angle_draw: T = self.unitFloatRound(T, index, 5);
const radius: T = @sqrt(-2.0 * @log(radius_draw));
return radius * @cos(two_pi * angle_draw);
}
/// An integer in `[0, bound)`, drawn at `index` from its own stream. `bound == 0` yields 0.
///
/// Lemire's multiply-shift, not `% bound`: the modulo is biased toward small values whenever
/// `bound` does not divide 2^32, and the bias stays invisible until something like a replay
/// buffer over-samples its early entries.
pub fn intBelow(self: Rng, index: u32, bound: u32) u32 {
if (bound == 0) {
return 0;
}
const wide: u64 = @as(u64, self.word(index, 7)) * @as(u64, bound);
return @intCast(wide >> 32);
}
/// Fill `out` with uniforms in `[0, 1)`, drawn at indices `0..out.len`.
pub fn fillUniform(self: Rng, comptime T: type, out: []T) void {
comptime requireFloat(T);
for (out, 0..) |*slot, i| {
slot.* = self.unitFloat(T, @intCast(i));
}
}
/// Fill `out` with standard normal draws, at indices `0..out.len`.
pub fn fillNormal(self: Rng, comptime T: type, out: []T) void {
comptime requireFloat(T);
for (out, 0..) |*slot, i| {
slot.* = self.normal(T, @intCast(i));
}
}
};
Tensor/// A view of `T`-typed storage as a multidimensional array.
///
/// A TENSOR NEVER OWNS ITS STORAGE, AND THAT IS THE DESIGN
///
/// There is no ownership flag and no `deinit`. `data` is a slice the tensor addresses; whoever
/// allocated it frees it, with the allocator that allocated it. Usually that is an arena, in
/// which case nobody frees anything and `arena.deinit()` ends the matter.
///
/// The alternative - an `owns: bool` and a `deinit` that is a no-op for views - is the shape
/// this library was warned about before it had a tensor: both of its failure modes are SILENT.
/// `deinit` on a view does nothing and reads as if it did something; a missing `deinit` on an
/// owner just grows memory. Having no `deinit` at all means a free appears exactly where an
/// allocation appears, and a view has neither.
///
/// Strides are in ELEMENTS and signed, so a reversed or negatively-strided view is expressible
/// without changing this type. `at` computes a signed offset and casts once.
pub fn Tensor(comptime T: type) type {
return struct {
const Self = @This();
/// The storage this tensor addresses. Not owned - see the type's doc comment.
data: []T,
/// Extent of each axis. Only the first `rank` entries are meaningful.
shape: [max_rank]usize,
/// Step in ELEMENTS between consecutive positions along each axis.
strides: [max_rank]isize,
rank: u8,
/// WHERE THIS VIEW STARTS INSIDE `data`
///
/// Every view until now began at element 0, which is why `slice` could not exist: a
/// window into the middle of a tensor IS an offset. With this, `slice` and the rest of
/// znum's `index` namespace become views rather than copies.
///
/// IT ALSO WEAKENS WHAT `data.len` MEANS. `data` is the whole backing store; `base`
/// and the strides say which of it this view addresses. Everything that walked
/// `data[0..size()]` assuming those coincide had to be corrected - five sites, all of
/// them fast paths that skip the index walk, which is exactly where such an assumption
/// hides.
base: usize = 0,
/// A dense row-major tensor over `shape`, allocated from `gpa`.
///
/// The caller frees `result.data` with the same allocator - or allocates from an arena
/// and frees nothing.
pub fn alloc(gpa: Allocator, shape: []const usize) Error!Self {
if (shape.len > max_rank) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
const storage: []T = gpa.alloc(T, count) catch return Error.OutOfMemory;
return fromSlice(storage, shape);
}
/// A dense row-major view of `data` as `shape`. `data.len` must be exactly the product of
/// the extents - a slice that is merely LARGE ENOUGH is rejected, because the usual reason
/// for the mismatch is a shape that is wrong rather than storage that is generous.
pub fn fromSlice(data: []T, shape: []const usize) Error!Self {
if (shape.len > max_rank) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
if (count != data.len) {
return Error.ShapeMismatch;
}
var self: Self = .{
.data = data,
.shape = @splat(1),
.strides = @splat(0),
.rank = @intCast(shape.len),
.base = 0,
};
// Row-major: the last axis is contiguous, and each earlier stride is the product of
// the extents to its right.
var stride: isize = 1;
var axis: usize = shape.len;
while (axis > 0) {
axis -= 1;
self.shape[axis] = shape[axis];
self.strides[axis] = stride;
stride *= @intCast(shape[axis]);
}
return self;
}
/// The number of elements this tensor addresses.
pub fn size(self: Self) usize {
var count: usize = 1;
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
count *= self.shape[axis];
}
return count;
}
/// Whether the strides are dense row-major, so `data[0..size()]` is the tensor in order.
///
/// An operation that walks memory linearly must ask this rather than assume it: a view
/// produced by `transpose` addresses the same storage in a different order, and a linear
/// walk over it silently reads the wrong elements rather than failing.
pub fn isContiguous(self: Self) bool {
var expected: isize = 1;
var axis: usize = self.rank;
while (axis > 0) {
axis -= 1;
if (self.shape[axis] != 1 and self.strides[axis] != expected) {
return false;
}
expected *= @intCast(self.shape[axis]);
}
return true;
}
/// The offset into `data` of the element at `indices`.
pub fn flatIndex(self: Self, indices: []const usize) Error!usize {
if (indices.len != self.rank) {
return Error.ShapeMismatch;
}
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
if (i >= self.shape[axis]) {
return Error.OutOfRange;
}
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}
/// The element at `indices`.
pub fn at(self: Self, indices: []const usize) Error!T {
return self.data[try self.flatIndex(indices)];
}
/// The flat offset for a coordinate, with no error.
///
/// THE 45 `flatIndex(...) catch unreachable` SITES ARE ALL THIS SHAPE
///
/// const at: []const usize = walker[0..out.rank];
/// out.data[out.offsetOf(at)] = ...
///
/// The slice is `walker[0..rank]`, so `indices.len != rank` **cannot** happen - the
/// caller built it from the rank. And the coordinate came from a `Walk` over this very
/// shape, so the bounds cannot fail either. Two impossible checks, one `catch
/// unreachable`, on the hottest line in the library.
///
/// This computes the offset and returns it. The remaining safety is `data[offset]`,
/// which is where Zig puts it: a panic with a message in Debug and ReleaseSafe.
///
/// **Use it where the rank is an invariant you have already established** - after a
/// shape check, or inside a walk over the tensor's own shape. Use `flatIndex` where the
/// coordinate came from outside and might be wrong.
pub fn offsetOf(self: Self, indices: []const usize) usize {
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}
/// Read from a rank-2 tensor. **Does not return an error.**
///
/// WHY `at` HAS TO AND THIS DOES NOT
///
/// `at(indices: []const usize)` can fail for two reasons. One is a real bounds check.
/// The other is `indices.len != rank` - and that check exists ONLY because a slice hides
/// its arity from the type system. Every call site writes a literal `&.{ i, j }`; the
/// compiler could know it is two, and does not.
///
/// Naming the arity in the function removes that half entirely. What remains is the
/// bounds check, and **Zig already has a name for that**: it is what `data[idx]` does -
/// a panic with a message and a stack trace in Debug and ReleaseSafe, the same UB as any
/// other slice index in ReleaseFast.
///
/// Nobody writes `arr[i] catch unreachable`. This file had **105 `catch unreachable`**,
/// most of them on an `at` whose rank the author had checked three lines earlier, and
/// that noise is the single biggest thing that makes numeric Zig look worse than it is.
///
/// The safety is not weaker - it is Zig's, in the place Zig puts it.
pub fn at2(self: Self, i: usize, j: usize) T {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i)) +
self.strides[1] * @as(isize, @intCast(j));
return self.data[@intCast(offset)];
}
/// Write to a rank-2 tensor. The mirror of `at2`, and error-free for the same reason.
pub fn setAt2(self: Self, i: usize, j: usize, value: T) void {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i)) +
self.strides[1] * @as(isize, @intCast(j));
self.data[@intCast(offset)] = value;
}
/// Read from a rank-1 tensor, error-free.
pub fn at1(self: Self, i: usize) T {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i));
return self.data[@intCast(offset)];
}
/// Write to a rank-1 tensor, error-free.
pub fn setAt1(self: Self, i: usize, value: T) void {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i));
self.data[@intCast(offset)] = value;
}
/// Write `value` at `indices`.
pub fn setAt(self: Self, indices: []const usize, value: T) Error!void {
self.data[try self.flatIndex(indices)] = value;
}
/// Set every element to `value`.
pub fn fill(self: Self, value: T) void {
if (self.isContiguous()) {
@memset(self.data[self.base..][0..self.size()], value);
return;
}
var walker: [max_rank]usize = @splat(0);
var remaining: usize = self.size();
while (remaining > 0) : (remaining -= 1) {
self.data[self.offsetOf(walker[0..self.rank])] = value;
advance(self.shape[0..self.rank], walker[0..self.rank]);
}
}
/// The same storage seen as `shape`. Requires a contiguous tensor and an equal element
/// count; both are checked rather than assumed.
pub fn reshape(self: Self, shape: []const usize) Error!Self {
if (!self.isContiguous()) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
if (count != self.size()) {
return Error.ShapeMismatch;
}
// From `base`, not 0: reshaping a slice must keep addressing that slice.
return fromSlice(self.data[self.base..][0..self.size()], shape);
}
/// Whether two distinct index positions address the same element, which happens exactly
/// when a stretched axis (stride 0) has extent above 1.
///
/// An operation that WRITES must refuse an aliased destination. Many indices landing on
/// one element means the last write wins and every earlier one vanishes - a result that is
/// wrong in a way no bounds check catches. `map` and `zip` check this so a caller who
/// passes a broadcast view as an output gets an error rather than a plausible answer.
pub fn isAliased(self: Self) bool {
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
if (self.shape[axis] > 1 and self.strides[axis] == 0) {
return true;
}
}
return false;
}
/// A read-only view of this tensor stretched to `shape` by the broadcasting rule.
///
/// BROADCASTING IS A STRIDE OF ZERO, NOT A COPY
///
/// An axis whose extent is 1 is stretched by setting its stride to 0, so every position
/// along it addresses the same element. No storage is allocated and no data moves; a
/// `(3, 1, 5)` tensor viewed as `(3, 4, 5)` is the same 15 values read four times.
///
/// **WRITING THROUGH A STRETCHED VIEW IS ALMOST CERTAINLY A MISTAKE.** Many indices
/// alias one element, so the last write wins and the others vanish. This returns a view
/// because that is what the type system has; treat it as an input.
///
/// `isContiguous` reports false for any stretched axis with extent > 1, which is what
/// stops a linear walk from reading the underlying 15 values as if they were 60.
pub fn broadcastTo(self: Self, shape: []const usize) Error!Self {
if (shape.len > max_rank or shape.len < self.rank) {
return Error.ShapeMismatch;
}
var view: Self = .{
.data = self.data,
.shape = @splat(1),
.strides = @splat(0),
.rank = @intCast(shape.len),
// Stretching does not move where the view starts.
.base = self.base,
};
// Align at the TRAILING axis, as the broadcast rule does: an axis this tensor does not
// have behaves as extent 1, and so does an axis it has with extent 1. Both get a
// stride of 0.
for (shape, 0..) |extent, axis| {
const from_end: usize = shape.len - axis;
view.shape[axis] = extent;
if (self.rank < from_end) {
continue;
}
const own_axis: usize = self.rank - from_end;
const own_extent: usize = self.shape[own_axis];
if (own_extent == extent) {
view.strides[axis] = self.strides[own_axis];
} else if (own_extent != 1) {
return Error.ShapeMismatch;
}
}
return view;
}
/// A window into this tensor: `len` positions along `axis`, starting at `start`.
///
/// A SLICE IS AN OFFSET, WHICH IS WHY `base` EXISTS
///
/// Nothing is copied: the view keeps every stride and simply starts `start` steps further
/// into the data along one axis. A batch of 32 rows out of 4096 is a `Tensor` the same
/// size as any other, addressing 32 rows of the same buffer.
///
/// `start + len` must fit. Clamping instead would hand back a view SMALLER than asked
/// for, and the caller's next loop would read the shape it requested rather than the one
/// it got.
pub fn slice(self: Self, axis: usize, start: usize, len: usize) Error!Self {
if (axis >= self.rank) {
return Error.OutOfRange;
}
if (start + len > self.shape[axis]) {
return Error.OutOfRange;
}
var view: Self = self;
const step: isize = self.strides[axis] * @as(isize, @intCast(start));
const moved: isize = @as(isize, @intCast(self.base)) + step;
if (moved < 0) {
return Error.OutOfRange;
}
view.base = @intCast(moved);
view.shape[axis] = len;
return view;
}
/// A view with the axes reordered: `out` axis `i` is this tensor's axis `order[i]`.
///
/// EVERY `manip` OPERATION HERE IS METADATA
///
/// `permute`, `squeeze`, `unsqueeze`, `flatten` and `moveAxis` all return a view: they
/// permute or insert entries in two fixed-size arrays and copy nothing. That is why they
/// need no GPU kernel and no sweep row - **there is no arithmetic to disagree about.**
/// The cost appears later, when a permuted view has to be BOUND to a kernel and must be
/// densified first, which is the constraint recorded in section 10.3.
///
/// `order` must be a permutation of `0..rank`. A repeated axis would produce a view
/// where two indices alias, which is `isAliased`'s job to reject at write time - but
/// catching it here names the actual mistake.
pub fn permute(self: Self, order: []const usize) Error!Self {
if (order.len != self.rank) {
return Error.ShapeMismatch;
}
var seen: [max_rank]bool = @splat(false);
for (order) |axis| {
if (axis >= self.rank or seen[axis]) {
return Error.OutOfRange;
}
seen[axis] = true;
}
var view: Self = self;
for (order, 0..) |axis, i| {
view.shape[i] = self.shape[axis];
view.strides[i] = self.strides[axis];
}
return view;
}
/// A view with a size-1 axis removed.
///
/// Removing an axis whose extent is NOT 1 would change how many elements the view
/// addresses, so it is `ShapeMismatch` rather than a silent reshape.
pub fn squeeze(self: Self, axis: usize) Error!Self {
if (axis >= self.rank) {
return Error.OutOfRange;
}
if (self.shape[axis] != 1) {
return Error.ShapeMismatch;
}
var view: Self = self;
var i: usize = axis;
while (i + 1 < self.rank) : (i += 1) {
view.shape[i] = self.shape[i + 1];
view.strides[i] = self.strides[i + 1];
}
view.shape[self.rank - 1] = 1;
view.strides[self.rank - 1] = 0;
view.rank = @intCast(self.rank - 1);
return view;
}
/// A view with a size-1 axis inserted at `axis`. `axis == rank` appends.
///
/// The new axis gets a stride of 0, which is harmless because its extent is 1 - a
/// single position cannot step anywhere. `isAliased` only objects when extent exceeds 1.
pub fn unsqueeze(self: Self, axis: usize) Error!Self {
if (axis > self.rank or self.rank + 1 > max_rank) {
return Error.OutOfRange;
}
var view: Self = self;
var i: usize = self.rank;
while (i > axis) : (i -= 1) {
view.shape[i] = self.shape[i - 1];
view.strides[i] = self.strides[i - 1];
}
view.shape[axis] = 1;
view.strides[axis] = 0;
view.rank = @intCast(self.rank + 1);
return view;
}
/// A rank-1 view of the whole tensor. Contiguous only, for `reshape`'s reason.
pub fn flatten(self: Self) Error!Self {
return self.reshape(&.{self.size()});
}
/// A view with one axis moved from `from` to `to`, the others keeping their order.
pub fn moveAxis(self: Self, from: usize, to: usize) Error!Self {
if (from >= self.rank or to >= self.rank) {
return Error.OutOfRange;
}
var order: [max_rank]usize = @splat(0);
var k: usize = 0;
var i: usize = 0;
while (i < self.rank) : (i += 1) {
if (i == from) {
continue;
}
order[k] = i;
k += 1;
}
// Open a slot at `to` and drop the moved axis into it.
var j: usize = self.rank - 1;
while (j > to) : (j -= 1) {
order[j] = order[j - 1];
}
order[to] = from;
return self.permute(order[0..self.rank]);
}
/// A view with axes `a` and `b` exchanged. Shares storage; changes through one are visible
/// through the other.
pub fn transpose(self: Self, a: usize, b: usize) Error!Self {
if (a >= self.rank or b >= self.rank) {
return Error.OutOfRange;
}
var view: Self = self;
view.shape[a] = self.shape[b];
view.shape[b] = self.shape[a];
view.strides[a] = self.strides[b];
view.strides[b] = self.strides[a];
return view;
}
};
}
linspace/// `count` evenly spaced values from `start` to `stop`.
///
/// `endpoint` is required for the reason `Ddof` and `CodeOrder` are: both conventions are in
/// common use, and picking the wrong one gives an answer that is off by one element everywhere
/// without being obviously wrong anywhere.
pub fn linspace(
comptime T: type,
gpa: Allocator,
start: T,
stop: T,
count: usize,
endpoint: EndPoint,
) Error!Tensor(T) {
comptime requireFloat(T);
const out: Tensor(T) = try Tensor(T).alloc(gpa, &.{count});
if (count == 0) {
return out;
}
if (count == 1) {
out.data[0] = start;
return out;
}
const divisor: T = @floatFromInt(if (endpoint == .include) count - 1 else count);
const step: T = (stop - start) / divisor;
for (0..count) |i| {
out.data[i] = start + @as(T, @floatFromInt(i)) * step;
}
// THE LAST VALUE IS SET EXACTLY, not computed. `start + (n-1) * step` is `stop` in real
// arithmetic and a few ulps off in floating point, and a caller who asked for a range ending
// at 1.0 and got 0.9999999999999999 has a bug report rather than a rounding lesson.
if (endpoint == .include) {
out.data[count - 1] = stop;
}
return out;
}
Densepub fn Dense(comptime T: type) type {
comptime requireFloat(T);
return struct {
const Self = @This();
weight: Tensor(T),
bias: Tensor(T),
/// Allocate and initialise. `He` for a layer feeding a rectifier, `Xavier` otherwise.
pub fn init(
gpa: Allocator,
rng: Rng,
fan_in: usize,
fan_out: usize,
scheme: InitScheme,
) Error!Self {
const weight: Tensor(T) = try Tensor(T).alloc(gpa, &.{ fan_in, fan_out });
const bias: Tensor(T) = try Tensor(T).alloc(gpa, &.{ 1, fan_out });
switch (scheme) {
.xavier => try initXavier(T, weight, rng, fan_in, fan_out),
.he => try initHe(T, weight, rng, fan_in),
}
// Biases start at zero. A random bias adds nothing a random weight does not already
// provide, and a non-zero one shifts every unit's activation before any data is seen.
bias.fill(0);
return .{ .weight = weight, .bias = bias };
}
/// What `attach` hands back: the output, and handles for the two parameters.
///
/// RETURNING ONLY THE OUTPUT IS NOT ENOUGH, which a first attempt at a training loop
/// showed immediately: the loop needs the parameter handles to read their gradients, and
/// with the layer holding no `Var` there is nowhere else to get them. Returning all three
/// keeps the layer graph-agnostic AND gives the caller what a step requires.
pub const Attached = struct {
out: Var,
weight: Var,
bias: Var,
};
/// Record `input @ weight + bias` onto `graph`.
pub fn attach(self: Self, graph: *Graph(T), input: Var) Error!Attached {
const w: Var = try graph.parameter(self.weight);
const b: Var = try graph.parameter(self.bias);
return .{
.out = try graph.add(try graph.matmul(input, w), b),
.weight = w,
.bias = b,
};
}
};
}
Graph/// Reverse-mode automatic differentiation over a recorded tape.
///
/// EAGER FORWARD, RECORDED BACKWARD
///
/// Each operation computes its result immediately AND appends a node describing how to propagate
/// a gradient back through it. Computing eagerly means every shape is known when it is needed, so
/// the graph allocates outputs as it goes and no separate shape-inference pass exists to disagree
/// with the forward pass.
///
/// THE TAPE IS A LIST, NOT A GRAPH SEARCH. Nodes are appended in execution order, and a value
/// can only depend on values recorded before it, so walking the list backwards visits every
/// consumer of a value before the value itself. That is a topological order by construction -
/// no sort, no visited set, no cycle check.
///
/// Everything is allocated from one allocator, usually an arena. A graph has no
/// `deinit`: closing the scope releases the tape, the values and the gradients together, which is
/// the same ownership rule the rest of the library follows.
pub fn Graph(comptime T: type) type {
comptime requireFloat(T);
return struct {
const Self = @This();
/// What a node needs to push a gradient from its output to its inputs.
const Op = enum {
leaf,
add,
sub_op,
mul_op,
matmul,
tanh_op,
relu_op,
sigmoid_op,
mse,
cross_entropy,
ppo_clip,
softmax_op,
layernorm_op,
scale_op,
dropout_op,
transpose_op,
exp_op,
log_op,
softplus_op,
elu_op,
slice_op,
concat_op,
conv2d_op,
embedding_op,
};
const Node = struct {
op: Op,
/// Up to two inputs. Unused slots hold the node's own index and are ignored.
lhs: Var = 0,
rhs: Var = 0,
/// `mse` records its target here; nothing else uses it.
target: Var = 0,
/// `cross_entropy` records its class indices here. A slice rather than a `Var`
/// because labels are integers the caller already holds, not values the graph
/// differentiates - putting them on the tape would imply a gradient that does not
/// exist.
labels: []const usize = &.{},
/// `scale_op` keeps its factor here, `layernorm_op` its epsilon, `dropout_op` its
/// keep-probability's reciprocal. One float, three meanings, each used by one op.
scalar: T = 0,
/// `dropout_op` keeps its mask as a value of the graph, so `recompute` reuses it and
/// `checkGradient` sees the same function it is differentiating.
mask: Var = 0,
/// `conv2d_op` keeps `{ stride_y, stride_x, pad_y, pad_x }` here. A dedicated field
/// rather than borrowing `labels`, which means class indices everywhere else - four
/// words per node is cheap next to a reader wondering which meaning is in play.
geometry: [4]usize = @splat(0),
/// Constants a node needs in its BACKWARD that are not `Var`s.
///
/// `ppo_clip` keeps the old log-probabilities and the advantages here: both are
/// fixed for the step - the old policy by definition, the advantages because they
/// were estimated before it - so neither is a tape node, and neither can be
/// recomputed from the graph. The alternative was to overload `labels`, which is
/// `[]const usize` and says "class indices"; a slot whose name is a lie costs more
/// than a field.
constants: []const T = &.{},
};
gpa: Allocator,
values: std.ArrayList(Tensor(T)),
grads: std.ArrayList(Tensor(T)),
/// Whether a value accumulates a gradient. Constants do not, which keeps the backward
/// pass from allocating and filling buffers nobody reads.
wants_grad: std.ArrayList(bool),
nodes: std.ArrayList(Node),
pub fn init(gpa: Allocator) Self {
return .{
.gpa = gpa,
.values = .empty,
.grads = .empty,
.wants_grad = .empty,
.nodes = .empty,
};
}
fn push(self: *Self, value: Tensor(T), node: Node, wants: bool) Error!Var {
const shape: []const usize = value.shape[0..value.rank];
const grad: Tensor(T) = if (wants) try Tensor(T).alloc(self.gpa, shape) else value;
if (wants) {
grad.fill(0);
}
try self.values.append(self.gpa, value);
try self.grads.append(self.gpa, grad);
try self.wants_grad.append(self.gpa, wants);
try self.nodes.append(self.gpa, node);
return @intCast(self.values.items.len - 1);
}
/// A value the graph will differentiate with respect to.
pub fn parameter(self: *Self, value: Tensor(T)) Error!Var {
return self.push(value, .{ .op = .leaf }, true);
}
/// A value that is fixed: inputs, targets, anything not being learned.
pub fn constant(self: *Self, value: Tensor(T)) Error!Var {
return self.push(value, .{ .op = .leaf }, false);
}
/// The tensor behind a value.
pub fn valueOf(self: *const Self, v: Var) Tensor(T) {
return self.values.items[v];
}
/// The gradient accumulated into a value by `backward`.
pub fn gradOf(self: *const Self, v: Var) Error!Tensor(T) {
if (!self.wants_grad.items[v]) {
return Error.DomainError;
}
return self.grads.items[v];
}
/// `out = a + b`, with `b` broadcast to `a`'s shape.
pub fn add(self: *Self, a: Var, b: Var) Error!Var {
const va: Tensor(T) = self.values.items[a];
const shape: []const usize = va.shape[0..va.rank];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try zimrnum.add(T, out, va, self.values.items[b]);
const wants: bool = self.wants_grad.items[a] or self.wants_grad.items[b];
return self.push(out, .{ .op = .add, .lhs = a, .rhs = b }, wants);
}
/// `out = a @ b`, rank 2.
pub fn matmul(self: *Self, a: Var, b: Var) Error!Var {
const va: Tensor(T) = self.values.items[a];
const vb: Tensor(T) = self.values.items[b];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, &.{ va.shape[0], vb.shape[1] });
try zimrnum.matmul(T, out, va, vb);
const wants: bool = self.wants_grad.items[a] or self.wants_grad.items[b];
return self.push(out, .{ .op = .matmul, .lhs = a, .rhs = b }, wants);
}
/// `out = a - b`, elementwise, same shape.
pub fn sub(self: *Self, a: Var, b: Var) Error!Var {
return self.binary(.sub_op, a, b);
}
/// `out = a * b`, elementwise, same shape.
pub fn mul(self: *Self, a: Var, b: Var) Error!Var {
return self.binary(.mul_op, a, b);
}
fn binary(self: *Self, op: Op, a: Var, b: Var) Error!Var {
const va: Tensor(T) = self.values.items[a];
const shape: []const usize = va.shape[0..va.rank];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
const vb: Tensor(T) = self.values.items[b];
switch (op) {
.sub_op => try zimrnum.sub(T, out, va, vb),
.mul_op => try zimrnum.mul(T, out, va, vb),
else => unreachable,
}
const wants: bool = self.wants_grad.items[a] or self.wants_grad.items[b];
return self.push(out, .{ .op = op, .lhs = a, .rhs = b }, wants);
}
/// `out = relu(a)`.
pub fn relu(self: *Self, a: Var) Error!Var {
return self.unary(.relu_op, a);
}
/// `out = sigmoid(a)`.
pub fn sigmoid(self: *Self, a: Var) Error!Var {
return self.unary(.sigmoid_op, a);
}
fn unary(self: *Self, op: Op, a: Var) Error!Var {
const va: Tensor(T) = self.values.items[a];
const shape: []const usize = va.shape[0..va.rank];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
switch (op) {
.relu_op => try zimrnum.relu(T, out, va),
.sigmoid_op => try zimrnum.sigmoid(T, out, va),
.tanh_op => try zimrnum.tanh(T, out, va),
else => unreachable,
}
return self.push(out, .{ .op = op, .lhs = a }, self.wants_grad.items[a]);
}
/// `out = tanh(a)`.
pub fn tanh(self: *Self, a: Var) Error!Var {
const va: Tensor(T) = self.values.items[a];
const shape: []const usize = va.shape[0..va.rank];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try zimrnum.tanh(T, out, va);
return self.push(out, .{ .op = .tanh_op, .lhs = a }, self.wants_grad.items[a]);
}
/// Softmax cross-entropy of rank-2 logits against class indices, as a 1x1 value.
///
/// The backward rule is `softmax(logits) - onehot`, over the row count - the pairing
/// `exp(x)`, elementwise.
///
/// THE LEVERAGED PAIR - AND WHY `softplus` STILL NEEDS ITS OWN NODE
///
/// `exp` and `log` unlock compositions the tape could not express: `softplus`, `mish`,
/// `logSumExp` as a differentiable reduction, anything with a ratio of exponentials.
/// That is the `slice`/`concat` shape again - the primitive is worth more than the
/// functions it is blocking.
///
/// But `softplus(x) = log(1 + exp(x))` **written that way is wrong**, and the third node
/// below exists because of it. Measured at f32: `exp(89)` is infinity, so the composition
/// returns `inf` from x = 89 upward where the true answer is 89 - softplus is very nearly
/// the identity for large x. **Having the primitives does not mean composing them is
/// safe**, which is worth knowing before reaching for them.
pub fn exp(self: *Self, a: Var) Error!Var {
const source: Tensor(T) = self.values.items[a];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, source.shape[0..source.rank]);
try zimrnum.exp(T, out, source);
return self.push(out, .{ .op = .exp_op, .lhs = a }, self.wants_grad.items[a]);
}
/// The natural logarithm, elementwise. A nonpositive input gives `-inf` or `nan`, which
/// is `@log`'s behaviour and is left alone: clamping here would hide the caller's bug.
pub fn log(self: *Self, a: Var) Error!Var {
const source: Tensor(T) = self.values.items[a];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, source.shape[0..source.rank]);
try zimrnum.log(T, out, source);
return self.push(out, .{ .op = .log_op, .lhs = a }, self.wants_grad.items[a]);
}
/// `log(1 + exp(x))`, computed so it neither overflows nor loses its tail.
///
/// ITS OWN NODE, BECAUSE THE COMPOSITION IS NOT SAFE
///
/// `log(1 + exp(x))` overflows at f32 from x = 89 and at f64 from about 710, returning
/// `inf` where the answer is x. The tensor version uses `max(x, 0) + log1p(exp(-|x|))`,
/// which is exact at both ends - see its own comment for the identity and for what the
/// two earlier forms lost.
///
/// **The backward is the reason this is cheap**: `d/dx softplus(x)` is exactly
/// `sigmoid(x)` - no special case, no branch, and nothing to derive.
pub fn softplus(self: *Self, a: Var) Error!Var {
const source: Tensor(T) = self.values.items[a];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, source.shape[0..source.rank]);
try zimrnum.softplus(T, out, source);
return self.push(out, .{ .op = .softplus_op, .lhs = a }, self.wants_grad.items[a]);
}
/// ELU: `x` above zero, `alpha * (exp(x) - 1)` below.
///
/// A REAL NODE, BECAUSE A BRANCH IS NOT A COMPOSITION
///
/// `silu`, `gelu` and `mish` are all products and sums of nodes that already exist, so
/// they needed nothing. ELU is not: it switches on the sign, and there is no elementwise
/// `select` here to build that from. **This is the last activation without a gradient**,
/// and closing it means one node rather than a fifth composition.
///
/// The backward is `1` above zero and `alpha * exp(x)` below - which is `elu(x) + alpha`,
/// so it could read the output. It reads the INPUT instead: the two agree mathematically
/// and only the input says which side of the branch a value was on. At exactly zero both
/// give `1`, so the kink is continuous in the derivative for `alpha = 1` and the test
/// probes across it rather than around it.
pub fn elu(self: *Self, a: Var, alpha: T) Error!Var {
const source: Tensor(T) = self.values.items[a];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, source.shape[0..source.rank]);
try zimrnum.elu(T, out, source, alpha);
return self.push(
out,
.{ .op = .elu_op, .lhs = a, .scalar = alpha },
self.wants_grad.items[a],
);
}
/// Mish: `x * tanh(softplus(x))`.
///
/// A composition again, now that `softplus` is a node - which is what the pair was for.
pub fn mish(self: *Self, x: Var) Error!Var {
return self.mul(x, try self.tanh(try self.softplus(x)));
}
/// SiLU, also called swish: `x * sigmoid(x)`.
///
/// ACTIVATIONS THE TAPE ALREADY CONTAINS, AND THE ONES IT DOES NOT
///
/// `silu`, `gelu`, `softplus` and `elu` all exist at tensor level here and **none of them
/// had a tape node**, so a network using any of them could be evaluated and not trained.
/// That is a worse gap than a missing gradient function: the forward works, the loss
/// falls for the layers that do have nodes, and the activation simply never learns.
///
/// Two of the four are compositions of nodes that already exist, so they cost nothing and
/// their gradients are the tape's:
///
/// silu(x) = x * sigmoid(x)
/// gelu(x) = 0.5x * (1 + tanh(0.7978845608 * (x + 0.044715 x^3)))
///
/// The other two are not, and the plan records why rather than half-building them:
/// `softplus` needs `log1p(exp(x))` and the tape has neither log nor exp; `elu` branches
/// on the sign, which is not an elementwise composition of anything here.
pub fn silu(self: *Self, x: Var) Error!Var {
return self.mul(x, try self.sigmoid(x));
}
/// GELU, in the tanh form every transformer uses.
///
/// The exact definition is `x * Phi(x)` with `Phi` the normal CDF, which needs `erf`. The
/// tanh approximation is what PyTorch, JAX and every published transformer actually run -
/// it agrees to about 1e-3 and is the thing weights were trained against, so matching it
/// matters more than matching the definition.
pub fn gelu(self: *Self, x: Var) Error!Var {
// `x * sigmoid(2u)`, NOT `0.5x * (1 + tanh(u))` - THE SAME FUNCTION, DERIVED
//
// (1 + tanh u) / 2
// = (1 + (e^u - e^-u)/(e^u + e^-u)) / 2
// = (2 e^u / (e^u + e^-u)) / 2
// = e^u / (e^u + e^-u)
// = 1 / (1 + e^-2u)
// = sigmoid(2u)
//
// The first form needs a tensor of ones to add - **allocated on every call**, and on
// a graph rebuilt each training step that is a per-step allocation for a constant.
// The second needs none, and is two nodes shorter. zimrmath's own `gelu` already used
// this identity; the tape's did not until it was read.
//
// `sqrt(2/pi)` is the 0.7978845608 in the published formula, doubled here because the
// identity takes `2u`.
const cubed: Var = try self.mul(x, try self.mul(x, x));
const inner: Var = try self.add(x, try self.scale(cubed, 0.044715));
return self.mul(x, try self.sigmoid(try self.scale(inner, 1.5957691216057308)));
}
/// A gated linear unit: the first half of the columns, gated by the second.
///
/// NO BACKWARD - AND THAT IS THE POINT
///
/// znum writes `gatedForward` and then a hand-derived `gatedBackward`, walking
/// coordinates and applying the product rule by hand. Here the whole thing is
///
/// slice(x, 0, half) * sigmoid(slice(x, half, half))
///
/// and every one of those is already a tape node, so **the gradient is the tape's**.
/// There is no second derivation to keep in step with the first, and no chance of the
/// pair disagreeing under maintenance.
///
/// This is the third thing `graph.slice` and `graph.concat` paid for, after multi-head
/// attention and the batched cartpole twin. The ops were built because their absence had
/// shaped two designs; they keep shaping ones that come after.
///
/// The last axis must be even - half the columns carry the value and half the gate.
pub fn glu(self: *Self, x: Var) Error!Var {
return self.gateWith(x, .sigmoid);
}
/// SwiGLU: the same shape, gated by SiLU rather than sigmoid.
///
/// The gate in every recent transformer feed-forward. `silu(v)` is `v * sigmoid(v)`, so
/// this too is existing nodes and nothing new.
pub fn swiglu(self: *Self, x: Var) Error!Var {
return self.gateWith(x, .silu);
}
/// Which gate `glu` and `swiglu` apply.
pub const Gate = enum { sigmoid, silu };
/// The shared body. Splitting once and gating once, with the gate chosen at comptime.
fn gateWith(self: *Self, x: Var, comptime gate: Gate) Error!Var {
const value: Tensor(T) = self.values.items[x];
if (value.rank != 2) {
return Error.UnsupportedShape;
}
const width: usize = value.shape[1];
if (width % 2 != 0 or width == 0) {
return Error.UnsupportedShape;
}
const half: usize = width / 2;
const carried: Var = try self.slice(x, 1, 0, half);
const gating: Var = try self.slice(x, 1, half, half);
const opened: Var = switch (gate) {
.sigmoid => try self.sigmoid(gating),
// SiLU is `v * sigmoid(v)` - two nodes, no new op.
.silu => try self.mul(gating, try self.sigmoid(gating)),
};
return self.mul(carried, opened);
}
/// PPO's clipped objective, AS A LOSS - the version `ppoClipObjective` could not be.
///
/// WHY THIS EXISTS: THE SCALAR ONE COMPUTES PPO AND CANNOT TRAIN WITH IT
///
/// `ppoClipObjective` returns a `T`. It is correct, it is tested, and **it is not on the
/// tape** - so it reports what the objective is worth and cannot produce a gradient. The
/// library had PPO's arithmetic and not PPO's training step, and nothing said so until an
/// end-to-end test tried to learn with it.
///
/// THE CLIP IS THE WHOLE ALGORITHM, AND IT LIVES IN THE BACKWARD
///
/// The forward takes a minimum, which looks like a detail. The gradient is where it acts:
/// **when the clip is the smaller branch, the gradient is ZERO** - the update stops
/// pushing a ratio that has already moved too far. That is what stops a policy step from
/// destroying the policy, and it is the reason PPO works where plain policy gradient
/// diverges.
///
/// `logp_old` and `advantages` are constants of the step: the old policy is fixed by
/// definition and the advantages were estimated before it. Only `logp_new` carries a
/// gradient, which is why it is the only `Var`.
///
/// Returns the NEGATED mean, so it is a loss to descend rather than an objective to
/// ascend - every other `*Loss` here descends, and one that did not would be a trap.
pub fn ppoClipLoss(
self: *Self,
logp_new: Var,
logp_old: []const T,
advantages: []const T,
clip: T,
) Error!Var {
const fresh: Tensor(T) = self.values.items[logp_new];
if (fresh.size() != logp_old.len or logp_old.len != advantages.len) {
return Error.ShapeMismatch;
}
if (clip <= 0 or fresh.size() == 0) {
return Error.DomainError;
}
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, &.{1});
out.data[0] = -(try ppoClipObjective(T, fresh.data, logp_old, advantages, clip));
// `labels` is the node's spare slot for caller data; the two constant slices ride
// there as one allocation so the backward can recompute the branch.
const kept: []T = try self.gpa.alloc(T, logp_old.len * 2);
@memcpy(kept[0..logp_old.len], logp_old);
@memcpy(kept[logp_old.len..], advantages);
return self.push(
out,
.{ .op = .ppo_clip, .lhs = logp_new, .scalar = clip, .constants = kept },
self.wants_grad.items[logp_new],
);
}
/// that makes the two functions' derivatives compose into a subtraction.
pub fn crossEntropy(self: *Self, logits: Var, labels: []const usize) Error!Var {
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, &.{ 1, 1 });
out.data[0] = try crossEntropyRows(T, self.values.items[logits], labels);
return self.push(
out,
.{ .op = .cross_entropy, .lhs = logits, .labels = labels },
self.wants_grad.items[logits],
);
}
/// Row-wise softmax of a rank-2 value.
pub fn softmaxRows(self: *Self, a: Var) Error!Var {
const va: Tensor(T) = self.values.items[a];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, va.shape[0..va.rank]);
try zimrnum.softmaxRows(T, out, va);
return self.push(out, .{ .op = .softmax_op, .lhs = a }, self.wants_grad.items[a]);
}
/// Row-wise layer normalisation of a rank-2 value.
pub fn layerNormRows(self: *Self, a: Var, epsilon: T) Error!Var {
const va: Tensor(T) = self.values.items[a];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, va.shape[0..va.rank]);
try zimrnum.layerNormRows(T, out, va, epsilon);
return self.push(
out,
.{ .op = .layernorm_op, .lhs = a, .scalar = epsilon },
self.wants_grad.items[a],
);
}
/// 2-D convolution of a rank-2 value with a learned kernel.
///
/// THE BACKWARD PASS IS THE FORWARD LOOP WITH THE ACCUMULATION REVERSED
///
/// Forward, each output reads a window: `y[i][j] += x[i*s+a-p][j*s+b-p] * k[a][b]`. Both
/// gradients fall out of that one line by moving what is read and what is written:
///
/// dx[i*s+a-p][j*s+b-p] += g[i][j] * k[a][b]
/// dk[a][b] += g[i][j] * x[i*s+a-p][j*s+b-p]
///
/// So the backward rule is the SAME loop nest with the same bounds and the same
/// in-range test - which is why it is written as one pass producing both, rather than as
/// the "full correlation with a flipped kernel" the textbooks give. That form is correct
/// and is a second thing to get right; this one cannot disagree with the forward pass
/// about which taps are in range, because it asks the identical question.
pub fn conv2d(
self: *Self,
input: Var,
kernel: Var,
stride: [2]usize,
pad: [2]usize,
) Error!Var {
const x: Tensor(T) = self.values.items[input];
const k: Tensor(T) = self.values.items[kernel];
if (x.rank != 2 or k.rank != 2) {
return Error.UnsupportedShape;
}
if (stride[0] == 0 or stride[1] == 0) {
return Error.DomainError;
}
const out_h: usize = (x.shape[0] + 2 * pad[0] - k.shape[0]) / stride[0] + 1;
const out_w: usize = (x.shape[1] + 2 * pad[1] - k.shape[1]) / stride[1] + 1;
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, &.{ out_h, out_w });
try zimrnum.conv2d(T, out, x, k, stride, pad);
return self.push(
out,
.{
.op = .conv2d_op,
.lhs = input,
.rhs = kernel,
.geometry = .{ stride[0], stride[1], pad[0], pad[1] },
},
self.wants_grad.items[input] or self.wants_grad.items[kernel],
);
}
/// A contiguous run along one axis, as a differentiable operation.
///
/// THE VIEW ALREADY EXISTS; WHAT IS NEW IS THE GRADIENT
///
/// `Tensor.slice` gives a strided view for free and always has. What it cannot do is
/// carry a gradient, so anything that slices a computed value has been out of reach:
/// multi-head attention splitting a projection, an LSTM slicing four fused gates out of
/// one matmul. **Both designs were shaped by this absence** - each said "stated rather
/// than half-built" and moved on. znum has no differentiable slice either.
///
/// The backward is a scatter: the gradient lands in the rows the slice took and every
/// other row gets zero, because a row nobody read cannot have influenced anything.
pub fn slice(self: *Self, a: Var, axis: usize, start: usize, len: usize) Error!Var {
const source: Tensor(T) = self.values.items[a];
if (axis >= source.rank or start + len > source.shape[axis]) {
return Error.OutOfRange;
}
var shape: [max_rank]usize = source.shape;
shape[axis] = len;
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, shape[0..source.rank]);
try materialise(T, out, try source.slice(axis, start, len));
return self.push(
out,
.{ .op = .slice_op, .lhs = a, .geometry = .{ axis, start, len, 0 } },
self.wants_grad.items[a],
);
}
/// Two values joined along one axis, as a differentiable operation.
///
/// The backward is the slice's, read in the other direction: each input gets back the
/// part of the gradient that came from its own rows.
pub fn concat(self: *Self, a: Var, b: Var, axis: usize) Error!Var {
const first: Tensor(T) = self.values.items[a];
const second: Tensor(T) = self.values.items[b];
if (first.rank != second.rank or axis >= first.rank) {
return Error.UnsupportedShape;
}
var shape: [max_rank]usize = first.shape;
shape[axis] = first.shape[axis] + second.shape[axis];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, shape[0..first.rank]);
try zimrnum.concat(T, out, first, second, axis);
return self.push(
out,
.{
.op = .concat_op,
.lhs = a,
.rhs = b,
.geometry = .{ axis, first.shape[axis], 0, 0 },
},
self.wants_grad.items[a] or self.wants_grad.items[b],
);
}
/// Gather rows of a table by index: `out[i] = table[rows[i]]`.
///
/// `rows` reuses the node's `labels` field, and here that is not a borrow - an
/// embedding's indices ARE the same kind of thing as a classifier's labels: positions
/// into a first axis, with no gradient of their own.
///
/// The gradient is a SCATTER-ADD, and `+=` is the whole point: a token appearing three
/// times in a batch must accumulate three contributions into one row of the table. A
/// scatter that assigned would train on the last occurrence and silently discard the
/// rest.
pub fn embedding(self: *Self, table: Var, rows: []const usize) Error!Var {
const t: Tensor(T) = self.values.items[table];
if (t.rank != 2) {
return Error.UnsupportedShape;
}
for (rows) |r| {
if (r >= t.shape[0]) {
return Error.OutOfRange;
}
}
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, &.{ rows.len, t.shape[1] });
try gatherRows(T, out, t, rows);
return self.push(
out,
.{ .op = .embedding_op, .lhs = table, .labels = rows },
self.wants_grad.items[table],
);
}
/// The transpose of a rank-2 value, MATERIALISED.
///
/// A view would be free on the CPU, but a value of the graph must be a tensor a later
/// `matmul` can read and a later `recompute` can overwrite in place - and a transposed
/// view of another value's storage is neither. The copy is the price of the tape owning
/// every value it records. Its gradient is the transpose of the incoming one.
pub fn transpose(self: *Self, a: Var) Error!Var {
const va: Tensor(T) = self.values.items[a];
if (va.rank != 2) {
return Error.UnsupportedShape;
}
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, &.{ va.shape[1], va.shape[0] });
try materialise(T, out, try va.transpose(0, 1));
return self.push(out, .{ .op = .transpose_op, .lhs = a }, self.wants_grad.items[a]);
}
/// `out = a * factor`.
pub fn scale(self: *Self, a: Var, factor: T) Error!Var {
const va: Tensor(T) = self.values.items[a];
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, va.shape[0..va.rank]);
try zimrnum.scale(T, out, va, factor);
return self.push(
out,
.{ .op = .scale_op, .lhs = a, .scalar = factor },
self.wants_grad.items[a],
);
}
/// Inverted dropout: each element is kept with probability `1 - rate` and scaled by
/// `1 / (1 - rate)`, so the expected value is unchanged.
///
/// THE MASK IS DRAWN ONCE AND STORED
///
/// A dropout that redrew its mask on every `recompute` would be a different function each
/// time - `checkGradient` would compare a derivative against a finite difference of
/// something else and report noise. The mask is a value of the graph, fixed until
/// `resampleDropout` is called, which a training loop does once per step.
///
/// Inverted scaling means nothing changes at inference: dropping the operation from the
/// graph is exactly equivalent to `rate = 0`, with no `1 - rate` factor to remember.
pub fn dropout(self: *Self, a: Var, rate: T, rng: Rng) Error!Var {
if (rate < 0 or rate >= 1) {
return Error.DomainError;
}
const va: Tensor(T) = self.values.items[a];
const shape: []const usize = va.shape[0..va.rank];
const mask: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
const keep: T = 1 - rate;
fillDropoutMask(T, mask, rng, keep);
const mask_var: Var = try self.constant(mask);
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try zimrnum.mul(T, out, va, mask);
return self.push(
out,
.{ .op = .dropout_op, .lhs = a, .mask = mask_var, .scalar = keep },
self.wants_grad.items[a],
);
}
/// Draw fresh masks for every dropout node, from a new generator.
pub fn resampleDropout(self: *Self, rng: Rng) void {
for (self.nodes.items) |node| {
if (node.op == .dropout_op) {
fillDropoutMask(T, self.values.items[node.mask], rng, node.scalar);
}
}
}
/// Mean squared error against `target`, as a 1x1 value so it is a `Var` like any other.
pub fn mseLoss(self: *Self, a: Var, target: Var) Error!Var {
const out: Tensor(T) = try Tensor(T).alloc(self.gpa, &.{ 1, 1 });
out.data[0] = try zimrnum.mseLoss(T, self.values.items[a], self.values.items[target]);
return self.push(
out,
.{ .op = .mse, .lhs = a, .target = target },
self.wants_grad.items[a],
);
}
/// Propagate a gradient of 1 from `loss` back to every parameter.
///
/// THE SEED IS 1, AND THAT IS WHY `loss` MUST BE A SINGLE VALUE. Differentiating a
/// tensor with respect to parameters gives a Jacobian, not a gradient; reverse mode
/// answers "how does THIS ONE NUMBER change", so the entry point takes a scalar and
/// seeds it with its own derivative with respect to itself.
pub fn backward(self: *Self, loss: Var) Error!void {
if (self.values.items[loss].size() != 1) {
return Error.UnsupportedShape;
}
for (self.grads.items, self.wants_grad.items) |g, wants| {
if (wants) {
g.fill(0);
}
}
self.grads.items[loss].fill(1);
var i: usize = self.nodes.items.len;
while (i > 0) {
i -= 1;
if (!self.wants_grad.items[i]) {
continue;
}
const node: Node = self.nodes.items[i];
const g: Tensor(T) = self.grads.items[i];
switch (node.op) {
.leaf => {},
.add => {
try self.accumulate(node.lhs, g);
try self.accumulate(node.rhs, g);
},
.sub_op => {
try self.accumulate(node.lhs, g);
const shape: []const usize = g.shape[0..g.rank];
const negated: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try neg(T, negated, g);
try self.accumulate(node.rhs, negated);
},
.mul_op => {
// Each side's gradient is the OTHER side's value times the incoming one.
const shape: []const usize = g.shape[0..g.rank];
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try zimrnum.mul(T, scratch, g, self.values.items[node.rhs]);
try self.accumulate(node.lhs, scratch);
const other: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try zimrnum.mul(T, other, g, self.values.items[node.lhs]);
try self.accumulate(node.rhs, other);
},
.relu_op => {
// reluGrad takes the forward INPUT, which is what `lhs` holds.
const in: Tensor(T) = self.values.items[node.lhs];
const shape: []const usize = in.shape[0..in.rank];
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try reluGrad(T, scratch, in, g);
try self.accumulate(node.lhs, scratch);
},
.sigmoid_op => {
// sigmoidGrad takes the forward OUTPUT, which is this node's value.
const out: Tensor(T) = self.values.items[i];
const shape: []const usize = out.shape[0..out.rank];
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try sigmoidGrad(T, scratch, out, g);
try self.accumulate(node.lhs, scratch);
},
.cross_entropy => {
const in: Tensor(T) = self.values.items[node.lhs];
const shape: []const usize = in.shape[0..in.rank];
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try crossEntropyRowsGrad(T, scratch, in, node.labels);
try zimrnum.scale(T, scratch, scratch, g.data[g.base]);
try self.accumulate(node.lhs, scratch);
},
.ppo_clip => {
// THE CLIP ACTS HERE, NOT IN THE FORWARD
//
// Where the clipped branch is the smaller one, the objective no longer
// depends on the ratio - so **the gradient is exactly zero** and the
// update stops pushing a ratio that has already moved too far. That is
// what keeps a policy step from destroying the policy.
//
// Elsewhere the objective is `ratio * advantage` and `d ratio/d logp` is
// `ratio`, so the derivative is `ratio * advantage`. The mean divides by
// the count, and the loss is negated, so both appear here too.
const fresh: Tensor(T) = self.values.items[node.lhs];
const n: usize = fresh.size();
const old: []const T = node.constants[0..n];
const advantages: []const T = node.constants[n..];
const scratch: Tensor(T) = try Tensor(T).alloc(
self.gpa,
fresh.shape[0..fresh.rank],
);
const low: T = 1 - node.scalar;
const high: T = 1 + node.scalar;
const seed: T = g.data[g.base] / @as(T, @floatFromInt(n));
for (0..n) |k| {
const ratio: T = @exp(fresh.data[k] - old[k]);
const unclipped: T = ratio * advantages[k];
const clipped: T = scalarClamp(ratio, low, high) * advantages[k];
// A zero where the clip won: the objective is flat in `ratio` there.
const active: bool = unclipped <= clipped;
scratch.data[k] = if (active) -seed * ratio * advantages[k] else 0;
}
try self.accumulate(node.lhs, scratch);
},
.softmax_op => {
// dx = y * (g - sum(g * y)) per row: the Jacobian of a softmax applied
// to g, without forming it.
const y: Tensor(T) = self.values.items[i];
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, y.shape[0..y.rank]);
try softmaxBackward(T, scratch, y, g);
try self.accumulate(node.lhs, scratch);
},
.layernorm_op => {
const y: Tensor(T) = self.values.items[i];
const x: Tensor(T) = self.values.items[node.lhs];
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, y.shape[0..y.rank]);
try layerNormBackward(T, scratch, x, y, g, node.scalar);
try self.accumulate(node.lhs, scratch);
},
.conv2d_op => {
const x: Tensor(T) = self.values.items[node.lhs];
const k: Tensor(T) = self.values.items[node.rhs];
const dx: Tensor(T) = try Tensor(T).alloc(self.gpa, x.shape[0..2]);
const dk: Tensor(T) = try Tensor(T).alloc(self.gpa, k.shape[0..2]);
dx.fill(0);
dk.fill(0);
try conv2dBackward(T, dx, dk, x, k, g, node.geometry);
try self.accumulate(node.lhs, dx);
try self.accumulate(node.rhs, dk);
},
.embedding_op => {
const t: Tensor(T) = self.values.items[node.lhs];
const dt: Tensor(T) = try Tensor(T).alloc(self.gpa, t.shape[0..2]);
dt.fill(0);
try scatterAddRows(T, dt, g, node.labels);
try self.accumulate(node.lhs, dt);
},
.transpose_op => {
const in: Tensor(T) = self.values.items[node.lhs];
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, in.shape[0..in.rank]);
try materialise(T, scratch, try g.transpose(0, 1));
try self.accumulate(node.lhs, scratch);
},
.exp_op => {
// `d/dx exp(x)` is `exp(x)` - the value already computed, so the
// backward reads the forward's output rather than recomputing it.
const value: Tensor(T) = self.values.items[i];
const scratch: Tensor(T) = try Tensor(T).alloc(
self.gpa,
value.shape[0..value.rank],
);
try zimrnum.mul(T, scratch, g, value);
try self.accumulate(node.lhs, scratch);
},
.log_op => {
// `d/dx log(x)` is `1/x`, and `x` is the INPUT - the one case here where
// the output is not enough.
const source: Tensor(T) = self.values.items[node.lhs];
const scratch: Tensor(T) = try Tensor(T).alloc(
self.gpa,
source.shape[0..source.rank],
);
try zimrnum.div(T, scratch, g, source);
try self.accumulate(node.lhs, scratch);
},
.softplus_op => {
// `d/dx softplus(x)` is exactly `sigmoid(x)` - no branch and nothing to
// derive, which is what makes a dedicated node cheap despite the stable
// forward being more than one line.
const source: Tensor(T) = self.values.items[node.lhs];
const scratch: Tensor(T) = try Tensor(T).alloc(
self.gpa,
source.shape[0..source.rank],
);
try zimrnum.sigmoid(T, scratch, source);
try zimrnum.mul(T, scratch, g, scratch);
try self.accumulate(node.lhs, scratch);
},
.elu_op => {
// `1` above zero, `alpha * exp(x)` below - read from the INPUT, because
// only the input says which side of the branch a value was on.
const source: Tensor(T) = self.values.items[node.lhs];
const scratch: Tensor(T) = try Tensor(T).alloc(
self.gpa,
source.shape[0..source.rank],
);
var walk: Walk = .over(source.shape[0..source.rank]);
while (walk.next()) |at| {
const x: T = try source.at(at);
const slope: T = if (x > 0) 1 else node.scalar * @exp(x);
try scratch.setAt(at, (try g.at(at)) * slope);
}
try self.accumulate(node.lhs, scratch);
},
.slice_op => {
// A SCATTER: the gradient lands where the slice read, and every other
// row gets zero - a row nobody read influenced nothing.
const source: Tensor(T) = self.values.items[node.lhs];
const scratch: Tensor(T) = try Tensor(T).alloc(
self.gpa,
source.shape[0..source.rank],
);
scratch.fill(0);
const window: Tensor(T) = try scratch.slice(
node.geometry[0],
node.geometry[1],
node.geometry[2],
);
var walk: Walk = .over(g.shape[0..g.rank]);
while (walk.next()) |at| {
try window.setAt(at, try g.at(at));
}
try self.accumulate(node.lhs, scratch);
},
.concat_op => {
// The slice's backward read the other way: each input takes back the
// part of the gradient that came from its own rows.
const axis: usize = node.geometry[0];
const split: usize = node.geometry[1];
const left: Tensor(T) = self.values.items[node.lhs];
const right: Tensor(T) = self.values.items[node.rhs];
const left_grad: Tensor(T) = try Tensor(T).alloc(
self.gpa,
left.shape[0..left.rank],
);
try materialise(T, left_grad, try g.slice(axis, 0, split));
try self.accumulate(node.lhs, left_grad);
const right_grad: Tensor(T) = try Tensor(T).alloc(
self.gpa,
right.shape[0..right.rank],
);
try materialise(T, right_grad, try g.slice(axis, split, right.shape[axis]));
try self.accumulate(node.rhs, right_grad);
},
.scale_op => {
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, g.shape[0..g.rank]);
try zimrnum.scale(T, scratch, g, node.scalar);
try self.accumulate(node.lhs, scratch);
},
.dropout_op => {
// The same mask that scaled the forward pass scales the gradient.
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, g.shape[0..g.rank]);
try zimrnum.mul(T, scratch, g, self.values.items[node.mask]);
try self.accumulate(node.lhs, scratch);
},
.matmul => try self.matmulBackward(node, g),
.tanh_op => {
const out: Tensor(T) = self.values.items[i];
const shape: []const usize = out.shape[0..out.rank];
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try tanhGrad(T, scratch, out, g);
try self.accumulate(node.lhs, scratch);
},
.mse => {
const a: Tensor(T) = self.values.items[node.lhs];
const shape: []const usize = a.shape[0..a.rank];
const scratch: Tensor(T) = try Tensor(T).alloc(self.gpa, shape);
try zimrnum.sub(T, scratch, a, self.values.items[node.target]);
const n: T = @floatFromInt(a.size());
try zimrnum.scale(T, scratch, scratch, 2 * g.data[g.base] / n);
try self.accumulate(node.lhs, scratch);
},
}
}
}
/// Replay the recorded operations forward, refreshing every computed value.
///
/// WHAT MAKES THE TAPE A MODEL RATHER THAN A ONE-SHOT
///
/// Without this a graph is single-use: a training loop would have to rebuild it every
/// step, allocating the whole tape again for a forward pass that differs only in the
/// numbers. Replaying instead writes new values into the buffers already allocated.
///
/// It is also what makes `checkGradient` possible. Verifying a derivative numerically
/// means evaluating the SAME function at a perturbed input, and "the same function" is
/// exactly what the tape is a record of.
///
/// Leaves are skipped: their values are whatever the caller last wrote, which is how a
/// new batch of inputs enters an existing graph.
pub fn recompute(self: *Self) Error!void {
for (self.nodes.items, 0..) |node, i| {
const out: Tensor(T) = self.values.items[i];
switch (node.op) {
.leaf => {},
.add => try zimrnum.add(
T,
out,
self.values.items[node.lhs],
self.values.items[node.rhs],
),
.sub_op => try zimrnum.sub(
T,
out,
self.values.items[node.lhs],
self.values.items[node.rhs],
),
.mul_op => try zimrnum.mul(
T,
out,
self.values.items[node.lhs],
self.values.items[node.rhs],
),
.matmul => try zimrnum.matmul(
T,
out,
self.values.items[node.lhs],
self.values.items[node.rhs],
),
.tanh_op => try zimrnum.tanh(T, out, self.values.items[node.lhs]),
.relu_op => try zimrnum.relu(T, out, self.values.items[node.lhs]),
.sigmoid_op => try zimrnum.sigmoid(T, out, self.values.items[node.lhs]),
.softmax_op => try zimrnum.softmaxRows(T, out, self.values.items[node.lhs]),
.layernorm_op => try zimrnum.layerNormRows(
T,
out,
self.values.items[node.lhs],
node.scalar,
),
.conv2d_op => try zimrnum.conv2d(
T,
out,
self.values.items[node.lhs],
self.values.items[node.rhs],
.{ node.geometry[0], node.geometry[1] },
.{ node.geometry[2], node.geometry[3] },
),
.embedding_op => try gatherRows(
T,
out,
self.values.items[node.lhs],
node.labels,
),
.transpose_op => try materialise(
T,
out,
try self.values.items[node.lhs].transpose(0, 1),
),
.exp_op => try zimrnum.exp(T, out, self.values.items[node.lhs]),
.log_op => try zimrnum.log(T, out, self.values.items[node.lhs]),
.softplus_op => try zimrnum.softplus(T, out, self.values.items[node.lhs]),
.elu_op => try zimrnum.elu(T, out, self.values.items[node.lhs], node.scalar),
.slice_op => try materialise(T, out, try self.values.items[node.lhs].slice(
node.geometry[0],
node.geometry[1],
node.geometry[2],
)),
.concat_op => try zimrnum.concat(
T,
out,
self.values.items[node.lhs],
self.values.items[node.rhs],
node.geometry[0],
),
.scale_op => try zimrnum.scale(T, out, self.values.items[node.lhs], node.scalar),
.dropout_op => try zimrnum.mul(
T,
out,
self.values.items[node.lhs],
self.values.items[node.mask],
),
.cross_entropy => out.data[out.base] = try crossEntropyRows(
T,
self.values.items[node.lhs],
node.labels,
),
.ppo_clip => {
const fresh: Tensor(T) = self.values.items[node.lhs];
const half: usize = fresh.size();
out.data[out.base] = -(try ppoClipObjective(
T,
fresh.data,
node.constants[0..half],
node.constants[half..],
node.scalar,
));
},
.mse => out.data[out.base] = try zimrnum.mseLoss(
T,
self.values.items[node.lhs],
self.values.items[node.target],
),
}
}
}
/// The largest disagreement between the tape's gradient for `param` and a central finite
/// difference of `loss`, over every element of `param`.
///
/// THE ONE CHECK THAT DOES NOT REPEAT THE DERIVATION
///
/// A hand-written backward pass and a tape rule can both be wrong in the same way,
/// because the same person derived them. A finite difference does not know the
/// derivative at all - it only evaluates the function twice - so agreement between the
/// two is evidence about the derivative rather than about the deriver.
///
/// CENTRAL, NOT FORWARD. `(f(x+h) - f(x-h)) / 2h` has error O(h^2) where
/// `(f(x+h) - f(x)) / h` has O(h); at a usable `h` that is the difference between
/// agreeing to six digits and to three, which decides whether a real error is visible
/// above the noise.
///
/// Every value is restored and the graph recomputed before returning, so calling this
/// leaves the model exactly as it found it.
pub fn checkGradient(self: *Self, loss: Var, param: Var, h: T) Error!T {
comptime requireFloat(T);
if (h <= 0) {
return Error.DomainError;
}
try self.recompute();
try self.backward(loss);
const analytic: Tensor(T) = try self.gradOf(param);
const value: Tensor(T) = self.values.items[param];
const scratch: Tensor(T) = try Tensor(T).alloc(
self.gpa,
analytic.shape[0..analytic.rank],
);
try materialise(T, scratch, analytic);
var worst: T = 0;
var index: usize = 0;
while (index < value.size()) : (index += 1) {
const at: usize = value.base + index;
const original: T = value.data[at];
value.data[at] = original + h;
try self.recompute();
const up: T = self.values.items[loss].data[self.values.items[loss].base];
value.data[at] = original - h;
try self.recompute();
const down: T = self.values.items[loss].data[self.values.items[loss].base];
value.data[at] = original;
const numeric: T = (up - down) / (2 * h);
worst = @max(worst, @abs(numeric - scratch.data[scratch.base + index]));
}
try self.recompute();
return worst;
}
/// `grad[v] += contribution`, summing over any axis `v` was broadcast along.
///
/// THE SUM IS WHAT MAKES BROADCASTING DIFFERENTIABLE. A bias of shape (1, n) added to
/// (m, n) is used m times in the forward pass, so its gradient is the sum of m
/// contributions. Adding the raw gradient instead would under-count by a factor of m and
/// produce a bias that learns m times too slowly - a bug that leaves the network training,
/// just badly, which is the hardest kind to notice.
fn accumulate(self: *Self, v: Var, contribution: Tensor(T)) Error!void {
if (!self.wants_grad.items[v]) {
return;
}
const dest: Tensor(T) = self.grads.items[v];
if (dest.size() == contribution.size()) {
return zimrnum.add(T, dest, dest, contribution);
}
// SUM OVER EVERY BROADCAST AXIS, NOT JUST THE LEADING ONE
//
// The first version summed rows only, which is right for a `(1, width)` bias under a
// `(rows, width)` gradient - `Dense`'s case, and the only one that existed. A `(1, 1)`
// bias broadcasting over BOTH axes, which `Conv2d` has, took column 0 and **silently
// dropped the rest**: a gradient wrong by a factor of the width, reported as a
// number rather than an error. The finite-difference check caught it at 0.7 where
// the kernel and input were at 8e-12.
//
// An axis of the destination that is 1 where the contribution is larger is a
// broadcast axis and is summed; an axis that matches is walked. Any other pairing is
// a shape error, which is what the first version should have returned.
if (dest.rank != contribution.rank) {
return Error.ShapeMismatch;
}
var axis: usize = 0;
while (axis < dest.rank) : (axis += 1) {
const same: bool = dest.shape[axis] == contribution.shape[axis];
if (!same and dest.shape[axis] != 1) {
return Error.ShapeMismatch;
}
}
var walker: [max_rank]usize = @splat(0);
var target: [max_rank]usize = @splat(0);
const shape: []const usize = contribution.shape[0..contribution.rank];
var remaining: usize = contribution.size();
while (remaining > 0) : (remaining -= 1) {
axis = 0;
while (axis < dest.rank) : (axis += 1) {
target[axis] = if (dest.shape[axis] == 1) 0 else walker[axis];
}
const at: []const usize = target[0..dest.rank];
const slot: usize = dest.offsetOf(at);
dest.data[slot] += contribution.data[contribution.offsetOf(walker[0..contribution.rank])];
advance(shape, walker[0..contribution.rank]);
}
}
/// `dA = g @ BT` and `dB = AT @ g`.
fn matmulBackward(self: *Self, node: Node, g: Tensor(T)) Error!void {
const va: Tensor(T) = self.values.items[node.lhs];
const vb: Tensor(T) = self.values.items[node.rhs];
if (self.wants_grad.items[node.lhs]) {
const da: Tensor(T) = try Tensor(T).alloc(
self.gpa,
&.{ va.shape[0], va.shape[1] },
);
try zimrnum.matmul(T, da, g, try vb.transpose(0, 1));
try self.accumulate(node.lhs, da);
}
if (self.wants_grad.items[node.rhs]) {
const db: Tensor(T) = try Tensor(T).alloc(
self.gpa,
&.{ vb.shape[0], vb.shape[1] },
);
try zimrnum.matmul(T, db, try va.transpose(0, 1), g);
try self.accumulate(node.rhs, db);
}
}
};
}
Var/// A handle to a value inside a `Graph`. An index, not a pointer.
///
/// WHY VALUES ARE INDICES AND NOT TENSORS
///
/// The obvious design keys gradients off the tensor itself - its `data.ptr`, say. That breaks
/// immediately here, because a view SHARES `data.ptr` with its parent: `w` and
/// `w.transpose(0,1)` would be the same key while being different values in the graph. Worse, a
/// tensor is copied by value throughout this library, so there is no stable address to key on at
/// all.
///
/// An index into the graph's own arrays sidesteps both. The graph owns the correspondence between
/// a value, its tensor, and its gradient, and nothing about `Tensor` has to change to support it.
pub const Var = u32;
Running it prints:
loss 9.4206 -> 0.0070
slope 3.054 (true 3.0), intercept 1.955 (true 2.0)Five things in that program are worth naming, each covered in full later:
| Line | What it is | Section |
|---|---|---|
zn.linspace(f64, gpa, 0, 1, 32, .include) | a tensor, built | 4.6 |
zn.Dense(f64).init(...) | a layer, owning its parameters | 11.9 |
zn.Graph(f64).init(gpa) | a fresh tape, once per step | 17.2 |
graph.backward(loss) | gradients for everything reachable | 17.2 |
w.* -= 0.1 * g | the update, written out | 11.7 |
The tape is rebuilt every step. That is deliberate and cheap: it is an arena of nodes, and discarding it is one pointer reset. Reusing a tape across steps would accumulate the whole history of training.
The update is written by hand here to show what an optimiser does. zn.adamWStep and
its family replace those four lines once the model has more than two parameters.
Inside the zimr tree the module is named zn. A module that wants it declares the
import in build.zig the same way it declares zm.
const zn = @import("zn");source .zenv.sh
zig build test-fast -Dtest-filter=zn. # zimrnum's tests only
zig build test-fast # the whole shader-free tier
zig build check # lint, shader corpus, regression gatesDraw a thousand standard normal values and report their mean and variance.
const zn = @import("zn");
var sample: [1000]f64 = undefined;
const rng: zn.Rng = zn.Rng.init(20260904);
rng.fillNormal(f64, &sample);
const t = try zn.Tensor(f64).fromSlice(&sample, &.{1000});
const mean = try zn.meanAll(f64, t);
const spread = try zn.variance(f64, t, .sample);
std.debug.assert(zn.approxEqAbs(f64, 0.0, mean, 0.1));
std.debug.assert(zn.approxEqAbs(f64, 1.0, spread, 0.1));Rng/// A reproducible random source addressed by (seed, index) rather than by mutation.
///
/// WHY COUNTER-BASED, AND NOT A STATEFUL GENERATOR
///
/// A classic PRNG carries state and advances it, so drawing the i-th value means drawing the
/// first i. That is fine on one CPU thread and useless everywhere else this library is going:
///
/// * **A GPU kernel has no shared state.** Thread `i` must be able to compute ITS value without
/// any other thread having run. A counter-based generator is a pure function of the index, so
/// the kernel is `value(seed, thread_id)` and needs no sequencing at all.
/// * **CPU and GPU must agree EXACTLY.** The CPU path is this library's oracle (plan section 5 rung
/// 1); an oracle that produces a different random stream than the thing it checks is not an
/// oracle. Same function, same integers, same answer.
/// * **A failing case is reproducible from two numbers.** "seed 7, index 4096" is the whole
/// state, so a bad draw can be re-created in a unit test without replaying a sequence.
///
/// THE MIXER IS DELIBERATELY MODEST AND THE TESTS SAY SO. `mix32` is the murmur3 finalizer -
/// a bijection with good avalanche, applied to `index` folded with `seed`. It is NOT a
/// cryptographic generator and NOT claimed to pass a full statistical battery; what it is pinned
/// on is the properties this library actually relies on: determinism, range, stream independence,
/// and first/second moments. A stronger generator can replace the body without touching a single
/// call site, precisely because the interface is (seed, index) rather than a stream.
pub const Rng = struct {
seed: u32,
/// A generator for a given seed. Two `Rng`s with the same seed are the same function.
pub fn init(seed: u32) Rng {
return .{ .seed = seed };
}
/// A DERIVED generator, so one seed can feed several independent uses without the caller
/// inventing seeds. `parent.split(0)` and `parent.split(1)` are unrelated streams.
///
/// Mixing the label rather than adding it: `seed + 1` and `seed + 2` are adjacent inputs to
/// the mixer, and adjacency is exactly what a hash is asked not to preserve. Passing the
/// label through `mix32` first means the derived seeds are as far apart as any other pair.
pub fn split(self: Rng, label: u32) Rng {
return .{ .seed = mix32(self.seed ^ mix32(label ^ 0x9e3779b9)) };
}
/// The number of mantissa bits `T` can hold, including the implicit leading one. This is how
/// many random bits a draw in `[0, 1)` may use: fewer wastes range, MORE ROUNDS UP TO 1.0 and
/// breaks the half-open interval.
fn mantissaWidth(comptime T: type) u32 {
comptime requireFloat(T);
return switch (T) {
f16 => 11,
f32 => 24,
f64 => 53,
else => unreachable,
};
}
/// One 32-bit word of the generator's output at `index`, from independent stream `stream`.
///
/// EACH KIND OF DRAW OWNS A STREAM, SO KINDS DO NOT INTERFERE
///
/// Without this, every method would start from `bits(index)` and drawing two different KINDS
/// at one index would return values that determine each other - `intBelow(i, n)` would be a
/// function of `unitFloat(i)`. A caller has no reason to expect that and every reason to be
/// hurt by it, so the stream is part of the addressing rather than the caller's problem.
///
/// The stream index is folded into the MIXER, not into the index. Folding it into the index
/// (`bits(index ^ K)`) makes indices `i` and `i ^ K` draw from the same two words with their
/// roles swapped, so those pairs are not independent - which is exactly the defect this
/// replaced.
///
/// `stream == 0` adds nothing, so `bits` is unchanged by the introduction of rounds.
fn word(self: Rng, index: u32, stream: u32) u32 {
return mix32(mix32(index *% 0x9e3779b9) ^ self.seed ^ (stream *% 0x85ebca6b));
}
/// The raw 32-bit draw at `index`, stream 0. Distinct from every other kind of draw.
pub fn bits(self: Rng, index: u32) u32 {
return self.word(index, 0);
}
/// A float in `[0, 1)` from stream `stream` (and `stream + 1` when `T` needs more than 32 bits).
///
/// Exactly `mantissaWidth(T)` random bits, scaled by 2^-width. The largest result is
/// (2^width - 1) / 2^width, which is the largest `T` strictly below one - so the interval is
/// half-open for every width, including `f16`, where using 24 bits rounds to 1.0 instead.
fn unitFloatRound(self: Rng, comptime T: type, index: u32, comptime stream: u32) T {
comptime requireFloat(T);
const width: u32 = comptime mantissaWidth(T);
const first: u64 = self.word(index, stream);
const mantissa: u64 = if (comptime width <= 32)
first >> comptime @as(u6, @intCast(32 - width))
else
(first << comptime @as(u6, @intCast(width - 32))) |
(@as(u64, self.word(index, stream + 1)) >> comptime @as(u6, @intCast(64 - width)));
const numerator: T = @floatFromInt(mantissa);
const denominator: T = @floatFromInt(@as(u64, 1) << comptime @as(u6, @intCast(width)));
return numerator / denominator;
}
/// A float in `[0, 1)`, drawn at `index`.
pub fn unitFloat(self: Rng, comptime T: type, index: u32) T {
return self.unitFloatRound(T, index, 1);
}
/// A float in `[low, high)`, drawn at `index`.
pub fn uniform(
self: Rng,
comptime T: type,
index: u32,
low: T,
high: T,
) T {
comptime requireFloat(T);
return low + (high - low) * self.unitFloat(T, index);
}
/// A standard normal draw (mean 0, variance 1) at `index`, by the Box-Muller transform.
///
/// ITS TWO UNIFORMS COME FROM ITS OWN STREAMS (3-6), not from the uniform stream, so
/// `normal(i)` and `unitFloat(i)` are independent and a caller mixing them needs no care.
///
/// Box-Muller rather than a ziggurat because a ziggurat REJECTS: the number of uniforms it
/// consumes varies per sample, which would make the value at `index` depend on how many
/// rejections happened. That is not a pure function of the index and would not survive being
/// evaluated on an independent GPU thread. Box-Muller consumes exactly two, always.
///
/// The radius draw is taken from `(0, 1]` rather than `[0, 1)` because `@log(0)` is -inf.
pub fn normal(self: Rng, comptime T: type, index: u32) T {
comptime requireFloat(T);
const two_pi: T = 6.283185307179586;
const radius_draw: T = 1.0 - self.unitFloatRound(T, index, 3);
const angle_draw: T = self.unitFloatRound(T, index, 5);
const radius: T = @sqrt(-2.0 * @log(radius_draw));
return radius * @cos(two_pi * angle_draw);
}
/// An integer in `[0, bound)`, drawn at `index` from its own stream. `bound == 0` yields 0.
///
/// Lemire's multiply-shift, not `% bound`: the modulo is biased toward small values whenever
/// `bound` does not divide 2^32, and the bias stays invisible until something like a replay
/// buffer over-samples its early entries.
pub fn intBelow(self: Rng, index: u32, bound: u32) u32 {
if (bound == 0) {
return 0;
}
const wide: u64 = @as(u64, self.word(index, 7)) * @as(u64, bound);
return @intCast(wide >> 32);
}
/// Fill `out` with uniforms in `[0, 1)`, drawn at indices `0..out.len`.
pub fn fillUniform(self: Rng, comptime T: type, out: []T) void {
comptime requireFloat(T);
for (out, 0..) |*slot, i| {
slot.* = self.unitFloat(T, @intCast(i));
}
}
/// Fill `out` with standard normal draws, at indices `0..out.len`.
pub fn fillNormal(self: Rng, comptime T: type, out: []T) void {
comptime requireFloat(T);
for (out, 0..) |*slot, i| {
slot.* = self.normal(T, @intCast(i));
}
}
};
Tensor/// A view of `T`-typed storage as a multidimensional array.
///
/// A TENSOR NEVER OWNS ITS STORAGE, AND THAT IS THE DESIGN
///
/// There is no ownership flag and no `deinit`. `data` is a slice the tensor addresses; whoever
/// allocated it frees it, with the allocator that allocated it. Usually that is an arena, in
/// which case nobody frees anything and `arena.deinit()` ends the matter.
///
/// The alternative - an `owns: bool` and a `deinit` that is a no-op for views - is the shape
/// this library was warned about before it had a tensor: both of its failure modes are SILENT.
/// `deinit` on a view does nothing and reads as if it did something; a missing `deinit` on an
/// owner just grows memory. Having no `deinit` at all means a free appears exactly where an
/// allocation appears, and a view has neither.
///
/// Strides are in ELEMENTS and signed, so a reversed or negatively-strided view is expressible
/// without changing this type. `at` computes a signed offset and casts once.
pub fn Tensor(comptime T: type) type {
return struct {
const Self = @This();
/// The storage this tensor addresses. Not owned - see the type's doc comment.
data: []T,
/// Extent of each axis. Only the first `rank` entries are meaningful.
shape: [max_rank]usize,
/// Step in ELEMENTS between consecutive positions along each axis.
strides: [max_rank]isize,
rank: u8,
/// WHERE THIS VIEW STARTS INSIDE `data`
///
/// Every view until now began at element 0, which is why `slice` could not exist: a
/// window into the middle of a tensor IS an offset. With this, `slice` and the rest of
/// znum's `index` namespace become views rather than copies.
///
/// IT ALSO WEAKENS WHAT `data.len` MEANS. `data` is the whole backing store; `base`
/// and the strides say which of it this view addresses. Everything that walked
/// `data[0..size()]` assuming those coincide had to be corrected - five sites, all of
/// them fast paths that skip the index walk, which is exactly where such an assumption
/// hides.
base: usize = 0,
/// A dense row-major tensor over `shape`, allocated from `gpa`.
///
/// The caller frees `result.data` with the same allocator - or allocates from an arena
/// and frees nothing.
pub fn alloc(gpa: Allocator, shape: []const usize) Error!Self {
if (shape.len > max_rank) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
const storage: []T = gpa.alloc(T, count) catch return Error.OutOfMemory;
return fromSlice(storage, shape);
}
/// A dense row-major view of `data` as `shape`. `data.len` must be exactly the product of
/// the extents - a slice that is merely LARGE ENOUGH is rejected, because the usual reason
/// for the mismatch is a shape that is wrong rather than storage that is generous.
pub fn fromSlice(data: []T, shape: []const usize) Error!Self {
if (shape.len > max_rank) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
if (count != data.len) {
return Error.ShapeMismatch;
}
var self: Self = .{
.data = data,
.shape = @splat(1),
.strides = @splat(0),
.rank = @intCast(shape.len),
.base = 0,
};
// Row-major: the last axis is contiguous, and each earlier stride is the product of
// the extents to its right.
var stride: isize = 1;
var axis: usize = shape.len;
while (axis > 0) {
axis -= 1;
self.shape[axis] = shape[axis];
self.strides[axis] = stride;
stride *= @intCast(shape[axis]);
}
return self;
}
/// The number of elements this tensor addresses.
pub fn size(self: Self) usize {
var count: usize = 1;
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
count *= self.shape[axis];
}
return count;
}
/// Whether the strides are dense row-major, so `data[0..size()]` is the tensor in order.
///
/// An operation that walks memory linearly must ask this rather than assume it: a view
/// produced by `transpose` addresses the same storage in a different order, and a linear
/// walk over it silently reads the wrong elements rather than failing.
pub fn isContiguous(self: Self) bool {
var expected: isize = 1;
var axis: usize = self.rank;
while (axis > 0) {
axis -= 1;
if (self.shape[axis] != 1 and self.strides[axis] != expected) {
return false;
}
expected *= @intCast(self.shape[axis]);
}
return true;
}
/// The offset into `data` of the element at `indices`.
pub fn flatIndex(self: Self, indices: []const usize) Error!usize {
if (indices.len != self.rank) {
return Error.ShapeMismatch;
}
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
if (i >= self.shape[axis]) {
return Error.OutOfRange;
}
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}
/// The element at `indices`.
pub fn at(self: Self, indices: []const usize) Error!T {
return self.data[try self.flatIndex(indices)];
}
/// The flat offset for a coordinate, with no error.
///
/// THE 45 `flatIndex(...) catch unreachable` SITES ARE ALL THIS SHAPE
///
/// const at: []const usize = walker[0..out.rank];
/// out.data[out.offsetOf(at)] = ...
///
/// The slice is `walker[0..rank]`, so `indices.len != rank` **cannot** happen - the
/// caller built it from the rank. And the coordinate came from a `Walk` over this very
/// shape, so the bounds cannot fail either. Two impossible checks, one `catch
/// unreachable`, on the hottest line in the library.
///
/// This computes the offset and returns it. The remaining safety is `data[offset]`,
/// which is where Zig puts it: a panic with a message in Debug and ReleaseSafe.
///
/// **Use it where the rank is an invariant you have already established** - after a
/// shape check, or inside a walk over the tensor's own shape. Use `flatIndex` where the
/// coordinate came from outside and might be wrong.
pub fn offsetOf(self: Self, indices: []const usize) usize {
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}
/// Read from a rank-2 tensor. **Does not return an error.**
///
/// WHY `at` HAS TO AND THIS DOES NOT
///
/// `at(indices: []const usize)` can fail for two reasons. One is a real bounds check.
/// The other is `indices.len != rank` - and that check exists ONLY because a slice hides
/// its arity from the type system. Every call site writes a literal `&.{ i, j }`; the
/// compiler could know it is two, and does not.
///
/// Naming the arity in the function removes that half entirely. What remains is the
/// bounds check, and **Zig already has a name for that**: it is what `data[idx]` does -
/// a panic with a message and a stack trace in Debug and ReleaseSafe, the same UB as any
/// other slice index in ReleaseFast.
///
/// Nobody writes `arr[i] catch unreachable`. This file had **105 `catch unreachable`**,
/// most of them on an `at` whose rank the author had checked three lines earlier, and
/// that noise is the single biggest thing that makes numeric Zig look worse than it is.
///
/// The safety is not weaker - it is Zig's, in the place Zig puts it.
pub fn at2(self: Self, i: usize, j: usize) T {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i)) +
self.strides[1] * @as(isize, @intCast(j));
return self.data[@intCast(offset)];
}
/// Write to a rank-2 tensor. The mirror of `at2`, and error-free for the same reason.
pub fn setAt2(self: Self, i: usize, j: usize, value: T) void {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i)) +
self.strides[1] * @as(isize, @intCast(j));
self.data[@intCast(offset)] = value;
}
/// Read from a rank-1 tensor, error-free.
pub fn at1(self: Self, i: usize) T {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i));
return self.data[@intCast(offset)];
}
/// Write to a rank-1 tensor, error-free.
pub fn setAt1(self: Self, i: usize, value: T) void {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i));
self.data[@intCast(offset)] = value;
}
/// Write `value` at `indices`.
pub fn setAt(self: Self, indices: []const usize, value: T) Error!void {
self.data[try self.flatIndex(indices)] = value;
}
/// Set every element to `value`.
pub fn fill(self: Self, value: T) void {
if (self.isContiguous()) {
@memset(self.data[self.base..][0..self.size()], value);
return;
}
var walker: [max_rank]usize = @splat(0);
var remaining: usize = self.size();
while (remaining > 0) : (remaining -= 1) {
self.data[self.offsetOf(walker[0..self.rank])] = value;
advance(self.shape[0..self.rank], walker[0..self.rank]);
}
}
/// The same storage seen as `shape`. Requires a contiguous tensor and an equal element
/// count; both are checked rather than assumed.
pub fn reshape(self: Self, shape: []const usize) Error!Self {
if (!self.isContiguous()) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
if (count != self.size()) {
return Error.ShapeMismatch;
}
// From `base`, not 0: reshaping a slice must keep addressing that slice.
return fromSlice(self.data[self.base..][0..self.size()], shape);
}
/// Whether two distinct index positions address the same element, which happens exactly
/// when a stretched axis (stride 0) has extent above 1.
///
/// An operation that WRITES must refuse an aliased destination. Many indices landing on
/// one element means the last write wins and every earlier one vanishes - a result that is
/// wrong in a way no bounds check catches. `map` and `zip` check this so a caller who
/// passes a broadcast view as an output gets an error rather than a plausible answer.
pub fn isAliased(self: Self) bool {
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
if (self.shape[axis] > 1 and self.strides[axis] == 0) {
return true;
}
}
return false;
}
/// A read-only view of this tensor stretched to `shape` by the broadcasting rule.
///
/// BROADCASTING IS A STRIDE OF ZERO, NOT A COPY
///
/// An axis whose extent is 1 is stretched by setting its stride to 0, so every position
/// along it addresses the same element. No storage is allocated and no data moves; a
/// `(3, 1, 5)` tensor viewed as `(3, 4, 5)` is the same 15 values read four times.
///
/// **WRITING THROUGH A STRETCHED VIEW IS ALMOST CERTAINLY A MISTAKE.** Many indices
/// alias one element, so the last write wins and the others vanish. This returns a view
/// because that is what the type system has; treat it as an input.
///
/// `isContiguous` reports false for any stretched axis with extent > 1, which is what
/// stops a linear walk from reading the underlying 15 values as if they were 60.
pub fn broadcastTo(self: Self, shape: []const usize) Error!Self {
if (shape.len > max_rank or shape.len < self.rank) {
return Error.ShapeMismatch;
}
var view: Self = .{
.data = self.data,
.shape = @splat(1),
.strides = @splat(0),
.rank = @intCast(shape.len),
// Stretching does not move where the view starts.
.base = self.base,
};
// Align at the TRAILING axis, as the broadcast rule does: an axis this tensor does not
// have behaves as extent 1, and so does an axis it has with extent 1. Both get a
// stride of 0.
for (shape, 0..) |extent, axis| {
const from_end: usize = shape.len - axis;
view.shape[axis] = extent;
if (self.rank < from_end) {
continue;
}
const own_axis: usize = self.rank - from_end;
const own_extent: usize = self.shape[own_axis];
if (own_extent == extent) {
view.strides[axis] = self.strides[own_axis];
} else if (own_extent != 1) {
return Error.ShapeMismatch;
}
}
return view;
}
/// A window into this tensor: `len` positions along `axis`, starting at `start`.
///
/// A SLICE IS AN OFFSET, WHICH IS WHY `base` EXISTS
///
/// Nothing is copied: the view keeps every stride and simply starts `start` steps further
/// into the data along one axis. A batch of 32 rows out of 4096 is a `Tensor` the same
/// size as any other, addressing 32 rows of the same buffer.
///
/// `start + len` must fit. Clamping instead would hand back a view SMALLER than asked
/// for, and the caller's next loop would read the shape it requested rather than the one
/// it got.
pub fn slice(self: Self, axis: usize, start: usize, len: usize) Error!Self {
if (axis >= self.rank) {
return Error.OutOfRange;
}
if (start + len > self.shape[axis]) {
return Error.OutOfRange;
}
var view: Self = self;
const step: isize = self.strides[axis] * @as(isize, @intCast(start));
const moved: isize = @as(isize, @intCast(self.base)) + step;
if (moved < 0) {
return Error.OutOfRange;
}
view.base = @intCast(moved);
view.shape[axis] = len;
return view;
}
/// A view with the axes reordered: `out` axis `i` is this tensor's axis `order[i]`.
///
/// EVERY `manip` OPERATION HERE IS METADATA
///
/// `permute`, `squeeze`, `unsqueeze`, `flatten` and `moveAxis` all return a view: they
/// permute or insert entries in two fixed-size arrays and copy nothing. That is why they
/// need no GPU kernel and no sweep row - **there is no arithmetic to disagree about.**
/// The cost appears later, when a permuted view has to be BOUND to a kernel and must be
/// densified first, which is the constraint recorded in section 10.3.
///
/// `order` must be a permutation of `0..rank`. A repeated axis would produce a view
/// where two indices alias, which is `isAliased`'s job to reject at write time - but
/// catching it here names the actual mistake.
pub fn permute(self: Self, order: []const usize) Error!Self {
if (order.len != self.rank) {
return Error.ShapeMismatch;
}
var seen: [max_rank]bool = @splat(false);
for (order) |axis| {
if (axis >= self.rank or seen[axis]) {
return Error.OutOfRange;
}
seen[axis] = true;
}
var view: Self = self;
for (order, 0..) |axis, i| {
view.shape[i] = self.shape[axis];
view.strides[i] = self.strides[axis];
}
return view;
}
/// A view with a size-1 axis removed.
///
/// Removing an axis whose extent is NOT 1 would change how many elements the view
/// addresses, so it is `ShapeMismatch` rather than a silent reshape.
pub fn squeeze(self: Self, axis: usize) Error!Self {
if (axis >= self.rank) {
return Error.OutOfRange;
}
if (self.shape[axis] != 1) {
return Error.ShapeMismatch;
}
var view: Self = self;
var i: usize = axis;
while (i + 1 < self.rank) : (i += 1) {
view.shape[i] = self.shape[i + 1];
view.strides[i] = self.strides[i + 1];
}
view.shape[self.rank - 1] = 1;
view.strides[self.rank - 1] = 0;
view.rank = @intCast(self.rank - 1);
return view;
}
/// A view with a size-1 axis inserted at `axis`. `axis == rank` appends.
///
/// The new axis gets a stride of 0, which is harmless because its extent is 1 - a
/// single position cannot step anywhere. `isAliased` only objects when extent exceeds 1.
pub fn unsqueeze(self: Self, axis: usize) Error!Self {
if (axis > self.rank or self.rank + 1 > max_rank) {
return Error.OutOfRange;
}
var view: Self = self;
var i: usize = self.rank;
while (i > axis) : (i -= 1) {
view.shape[i] = self.shape[i - 1];
view.strides[i] = self.strides[i - 1];
}
view.shape[axis] = 1;
view.strides[axis] = 0;
view.rank = @intCast(self.rank + 1);
return view;
}
/// A rank-1 view of the whole tensor. Contiguous only, for `reshape`'s reason.
pub fn flatten(self: Self) Error!Self {
return self.reshape(&.{self.size()});
}
/// A view with one axis moved from `from` to `to`, the others keeping their order.
pub fn moveAxis(self: Self, from: usize, to: usize) Error!Self {
if (from >= self.rank or to >= self.rank) {
return Error.OutOfRange;
}
var order: [max_rank]usize = @splat(0);
var k: usize = 0;
var i: usize = 0;
while (i < self.rank) : (i += 1) {
if (i == from) {
continue;
}
order[k] = i;
k += 1;
}
// Open a slot at `to` and drop the moved axis into it.
var j: usize = self.rank - 1;
while (j > to) : (j -= 1) {
order[j] = order[j - 1];
}
order[to] = from;
return self.permute(order[0..self.rank]);
}
/// A view with axes `a` and `b` exchanged. Shares storage; changes through one are visible
/// through the other.
pub fn transpose(self: Self, a: usize, b: usize) Error!Self {
if (a >= self.rank or b >= self.rank) {
return Error.OutOfRange;
}
var view: Self = self;
view.shape[a] = self.shape[b];
view.shape[b] = self.shape[a];
view.strides[a] = self.strides[b];
view.strides[b] = self.strides[a];
return view;
}
};
}
meanAll/// The arithmetic mean of every element, or `DomainError` for an empty tensor.
///
/// An empty mean is undefined rather than zero or NaN, and saying so is the caller's cue to
/// decide what an empty batch means - which is a question only the caller can answer.
pub fn meanAll(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n == 0) {
return Error.DomainError;
}
const divisor: T = @floatFromInt(n);
return sumAll(T, a) / divisor;
}
variance/// The variance of every element, with the divisor chosen by `ddof`.
///
/// TWO PASSES, FOR THE REASON `layerNormRows` GIVES
///
/// `E[x^2] - E[x]^2` is one pass and subtracts two large nearly-equal numbers, losing most of its
/// significant digits when the mean is large relative to the spread. Two passes cost one more
/// read of memory that is already in cache and are correct for every input.
///
/// Returns `DomainError` when the divisor is undefined - `n = 0`, or `n = 1` with
/// `.sample` - rather than a NaN that propagates into a result nobody examines. That is what
/// `Ddof.divisor` returning an optional is for.
pub fn variance(comptime T: type, a: Tensor(T), ddof: Ddof) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
const divisor: f64 = ddof.divisor(n) orelse return Error.DomainError;
const mean: T = try meanAll(T, a);
var total: CompensatedSum(T) = .{};
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
const d: T = (try a.at(at)) - mean;
const term: T = d * d;
total.add(term);
}
const divide_by: T = @floatCast(divisor);
return total.value() / divide_by;
}
approxEqAbs/// True when `a` and `b` agree to within `tol` in ABSOLUTE terms, with NaN never equal to
/// anything and infinities equal only to themselves.
///
/// Absolute, not relative, and that is a deliberate limitation: a relative tolerance is the
/// right tool near zero and the wrong one for a residual that is SUPPOSED to be near zero. A
/// caller comparing large magnitudes should scale its own tolerance and say so.
pub fn approxEqAbs(comptime T: type, a: T, b: T, tol: T) bool {
comptime requireFloat(T);
// `x != x` rather than a helper: `zm.isNan` is VECTOR-only (it returns
// `@Vector(veclen(T), bool)` and `veclen` rejects a scalar float). The inequality is
// the definition, is GPU-portable, and needs nothing imported.
if (a != a or b != b) {
return false;
}
if (a == b) {
return true;
}
return @abs(a - b) <= tol;
}
Every code block in this document is either the body of a test in
src/zimrnum.zig or built from one.
Every float-taking entry point is generic over f16, f32 and
f64; there is no default. Use f32 to match the engine and anything
destined for a GPU buffer. Use f64 where precision determines whether the answer
means anything: finite differences, long accumulations, ill-conditioned solves.
isFloat answers the question; requireFloat is the guard the library
calls, which fails at your call site with your type named rather than deeper down.
pub fn isFloat(comptime T: type) bool {
return switch (T) {
f16, f32, f64 => true,
else => false,
};
}zn.Error is one set for the whole library, so a caller handling a shape mismatch
does not need to know which layer noticed.
| Error | Raised when |
|---|---|
OutOfMemory | An allocation failed. |
ShapeMismatch | Operands disagree about shape and no broadcasting rule reconciles them; an output is sized wrong. |
OutOfRange | An axis, index or window is outside the tensor it addresses. |
DomainError | Defined, but not for these values: an empty mean, a variance with no degrees of freedom, a constant tensor to minMaxScale. |
UnsupportedShape | A shape the operation cannot accept: rank other than 2 to matmul, a non-contiguous tensor to reshape, an aliased destination to any writing operation. |
Whether a variance divides by n or n − 1. NumPy spells
this ddof=0 and ddof=1; the named form reads at the call site.
pub fn divisor(self: Ddof, n: usize) ?f64 {
return switch (self) {
.population => if (n == 0) null else float64(n),
.sample => if (n < 2) null else float64(n - 1),
};
}n | .population | .sample |
|---|---|---|
| 0 | null | null |
| 1 | 1.0 | null |
| 4 | 4.0 | 3.0 |
The optional is the point: a sample of one has a population variance — zero — and no
unbiased sample variance. variance turns the null into
DomainError rather than dividing by zero.
zn.approxEqAbs(f64, expected, actual, 1.0e-9)approxEqAbs/// True when `a` and `b` agree to within `tol` in ABSOLUTE terms, with NaN never equal to
/// anything and infinities equal only to themselves.
///
/// Absolute, not relative, and that is a deliberate limitation: a relative tolerance is the
/// right tool near zero and the wrong one for a residual that is SUPPOSED to be near zero. A
/// caller comparing large magnitudes should scale its own tolerance and say so.
pub fn approxEqAbs(comptime T: type, a: T, b: T, tol: T) bool {
comptime requireFloat(T);
// `x != x` rather than a helper: `zm.isNan` is VECTOR-only (it returns
// `@Vector(veclen(T), bool)` and `veclen` rejects a scalar float). The inequality is
// the definition, is GPU-portable, and needs nothing imported.
if (a != a or b != b) {
return false;
}
if (a == b) {
return true;
}
return @abs(a - b) <= tol;
}
Absolute tolerance. NaN is equal to nothing, including itself; infinities are
equal only to themselves. The tolerance is absolute and not relative, which suits a residual that
should be near zero. If you are comparing large magnitudes, scale the tolerance yourself.
Thirty-five tensor operations here are named after a zimrmath function and are meant to be that function at a higher rank. A comptime gate asserts the name exists on both sides, and fails the build if one is added without the other.
That is a spelling check. Two functions can share a name and disagree, and one
did: pow computed exp(y * log(x)), which is NaN for a negative base, while
zm.pow(-2, 3) returns −8. The doc comment called that deliberate. Nothing would
ever have reported it.
pow(-2, 3) -> NaN
zm.pow(-2, 3) -> -8Two tests now compare answers. One sweeps 42 base/power pairs for pow; the
other runs all thirty unary names over a probe that deliberately leaves several domains —
negatives for sqrt, values above one for atanh, zero for
log. NaN and infinity are accepted on both sides rather than skipped, because
those are exactly the cases where two implementations drift.
360 comparisons, zero further disagreements. That result is worth as much as the find: the pairing can be relied on, and the sweep is what keeps it that way.
A third test covers the four remaining names — atan2Rad, hypot,
clamp, lerp — which are the group pow came from. A
function taking two operands has more ways to disagree than a unary one:
atan2(y, x) takes the vertical component first,
and a version passing them the other way is wrong everywhere except the diagonal. The probe is
deliberately not symmetric, because a symmetric one would miss it.clamp is swept with an inverted range
(lo > hi). Whatever the two do there, they must do the same thing.lerp with t at −3 and
3 is extrapolation, not an error, and both must extrapolate alike.All thirty-five paired names are now checked by value, not just by name.
Every one of these tests was checked to fail — the unary sweep by perturbing an
implementation by 1e−7, the non-unary one by swapping atan2's arguments, which
two tests caught. A test that has never been seen to fail is not yet a test.
zn.zeros(f64, gpa, &.{ 2, 3 }) every element 0
zn.ones(f64, gpa, &.{4}) every element 1
zn.full(i32, gpa, &.{3}, 7) every element 7
zn.eye(f64, gpa, 3, 3) identity
zn.arange(i32, gpa, 0, 5, 1) 0 1 2 3 4
zn.linspace(f64, gpa, 0, 1, 5, .include)zeros/// A tensor of the given shape, every element zero.
///
/// `alloc` LEAVES GARBAGE, WHICH IS WHY THIS IS NOT THE SAME FUNCTION
///
/// Allocating and filling are separate steps and the allocator does not zero. A caller who
/// writes `alloc` and forgets the fill reads whatever was in that memory - plausible numbers,
/// sometimes, which is the worst kind.
///
/// So the common case gets a name. `alloc` stays for the case where every element is about to be
/// written anyway and zeroing would be wasted work.
pub fn zeros(comptime T: type, gpa: Allocator, shape: []const usize) Error!Tensor(T) {
comptime requireNumeric(T);
const out: Tensor(T) = try Tensor(T).alloc(gpa, shape);
out.fill(0);
return out;
}
ones/// A tensor of the given shape, every element one.
pub fn ones(comptime T: type, gpa: Allocator, shape: []const usize) Error!Tensor(T) {
comptime requireNumeric(T);
const out: Tensor(T) = try Tensor(T).alloc(gpa, shape);
out.fill(1);
return out;
}
full/// A tensor of the given shape, every element `value`.
pub fn full(
comptime T: type,
gpa: Allocator,
shape: []const usize,
value: T,
) Error!Tensor(T) {
comptime requireNumeric(T);
const out: Tensor(T) = try Tensor(T).alloc(gpa, shape);
out.fill(value);
return out;
}
eye/// An identity matrix, `rows` by `columns`.
///
/// RECTANGULAR IS ALLOWED, AND IT IS NOT A CURIOSITY
///
/// A non-square `eye` has ones down the main diagonal until one side runs out. That is what you
/// want for a projection between spaces of different sizes, and for the initial state of a
/// decomposition whose factors are not square. numpy allows it; refusing it would send a caller
/// to build one by hand.
pub fn eye(
comptime T: type,
gpa: Allocator,
rows: usize,
columns: usize,
) Error!Tensor(T) {
comptime requireNumeric(T);
const out: Tensor(T) = try zeros(T, gpa, &.{ rows, columns });
for (0..@min(rows, columns)) |i| {
out.setAt2(i, i, 1);
}
return out;
}
arange/// Evenly spaced values from `start`, stepping by `step`, stopping before `stop`.
///
/// THE ENDPOINT IS EXCLUDED, WHICH IS `arange`'S CONVENTION EVERYWHERE
///
/// `arange(0, 5, 1)` is `0 1 2 3 4` - five values, not six. numpy, Python's `range` and Zig's
/// `0..5` all agree, and a version that included the stop would be surprising in three languages
/// at once.
///
/// `linspace` is the other convention and takes a named argument, because there the count is
/// what you asked for and the endpoint is genuinely a choice.
pub fn arange(
comptime T: type,
gpa: Allocator,
start: T,
stop: T,
step: T,
) Error!Tensor(T) {
comptime requireNumeric(T);
if (step == 0) {
return Error.DomainError;
}
// A range that steps away from its stop has no end, and returning an empty tensor would hide
// the mistake. Counting first also means one allocation of exactly the right size.
var count: usize = 0;
if ((step > 0 and stop > start) or (step < 0 and stop < start)) {
const span: T = stop - start;
count = switch (@typeInfo(T)) {
.int => @intCast(@divTrunc(span + step - (if (step > 0) @as(T, 1) else @as(T, -1)), step)),
else => @intFromFloat(@ceil(span / step)),
};
}
const out: Tensor(T) = try Tensor(T).alloc(gpa, &.{count});
for (0..count) |i| {
// `start + i * step`, NOT a running sum: a running sum accumulates rounding, so the last
// value of `arange(0, 1, 0.1)` would drift. Multiplying keeps each value independent.
out.data[i] = start + @as(T, switch (@typeInfo(T)) {
.int => @intCast(i),
else => @floatFromInt(i),
}) * step;
}
return out;
}
linspace/// `count` evenly spaced values from `start` to `stop`.
///
/// `endpoint` is required for the reason `Ddof` and `CodeOrder` are: both conventions are in
/// common use, and picking the wrong one gives an answer that is off by one element everywhere
/// without being obviously wrong anywhere.
pub fn linspace(
comptime T: type,
gpa: Allocator,
start: T,
stop: T,
count: usize,
endpoint: EndPoint,
) Error!Tensor(T) {
comptime requireFloat(T);
const out: Tensor(T) = try Tensor(T).alloc(gpa, &.{count});
if (count == 0) {
return out;
}
if (count == 1) {
out.data[0] = start;
return out;
}
const divisor: T = @floatFromInt(if (endpoint == .include) count - 1 else count);
const step: T = (stop - start) / divisor;
for (0..count) |i| {
out.data[i] = start + @as(T, @floatFromInt(i)) * step;
}
// THE LAST VALUE IS SET EXACTLY, not computed. `start + (n-1) * step` is `stop` in real
// arithmetic and a few ulps off in floating point, and a caller who asked for a range ending
// at 1.0 and got 0.9999999999999999 has a bug report rather than a rounding lesson.
if (endpoint == .include) {
out.data[count - 1] = stop;
}
return out;
}
alloc is still there and still leaves garbage — the
allocator does not zero. A caller who allocates and forgets to fill reads whatever was in that
memory, which is sometimes plausible numbers. So the common case gets a name, and
alloc stays for when every element is about to be overwritten anyway.
arange(0, 5, 1) -> 0 1 2 3 4 five values, stop EXCLUDED
linspace(0, 1, 5, .include) -> 0 .25 .5 .75 1 five values, stop INCLUDED
linspace(0, 1, 5, .exclude) -> 0 .2 .4 .6 .8 five values, stop excludedarange excludes its stop and does not ask. numpy, Python's range and
Zig's 0..5 all agree, and a version that included it would be surprising in three
languages at once.
linspace asks, because there the count is what you requested and
the endpoint is a genuine choice. Same spacing question, different answers — and the
argument is required for the reason Ddof and CodeOrder are: the wrong one
gives a range that is off by one element everywhere without being obviously wrong anywhere.
start + (n-1) * step is stop in real arithmetic and a few ulps off in
floating point. A caller who asked for a range ending at 1.0 and got 0.9999999999999999 has a bug
report rather than a rounding lesson, so .include assigns the final slot
directly.
.exclude has no final value to pin, so every entry is
start + i * step and 3 * 0.2 comes out 0.6000000000000001. That is
correct arithmetic, not a defect — and it approximately where the
inclusive case is asserted exactly, which is the distinction in one place.
eye(2, 4) has ones down the main diagonal until one side runs out. That is what a
projection between spaces of different sizes needs, and what a decomposition whose factors are not
square starts from. numpy allows it; refusing it would send a caller to build one by hand.
try zn.astype(f64, i32, out, a); // f64 tensor -> i32 tensorastype/// Convert every element to another numeric type.
///
/// WHAT THIS HAS TO DECIDE, AND WHY numpy'S ANSWER DOES NOT TRANSFER
///
/// Four conversions and four different hazards:
///
/// f64 -> f32 1e300 becomes infinity
/// f64 -> i32 3.7 becomes 3 or 4, depending on a rounding rule
/// f64 -> i8 300.0 is OUT OF RANGE
/// i32 -> f32 2^24 + 1 loses its last bit
///
/// numpy truncates toward zero and **wraps out-of-range values silently**, which is how a
/// temperature of 300 becomes 44 in an i8 column and nobody notices for a week.
///
/// Zig cannot follow that even if it wanted to: `@intFromFloat` out of range is **undefined
/// behaviour**, not a wrap. So the choice is forced - check, and this returns `DomainError`.
///
/// TRUNCATION, NAMED
///
/// `3.7` becomes `3` and `-3.7` becomes `-3`: truncation toward zero, which is `@intFromFloat`'s
/// rule, C's, and numpy's. It is not `floor`, and the negative case is where they differ - `floor`
/// would give `-4`. Callers who want a different rounding apply `round` or `floor` FIRST, which
/// makes the choice visible at the call site instead of hidden in a conversion.
///
/// NaN AND INFINITY HAVE NO INTEGER, so they are a `DomainError` too rather than whatever the
/// hardware happens to produce.
pub fn astype(
comptime From: type,
comptime To: type,
out: Tensor(To),
a: Tensor(From),
) Error!void {
comptime requireNumeric(From);
comptime requireNumeric(To);
if (out.rank != a.rank) {
return Error.ShapeMismatch;
}
for (0..a.rank) |d| {
if (out.shape[d] != a.shape[d]) {
return Error.ShapeMismatch;
}
}
const to_int: bool = comptime @typeInfo(To) == .int;
const from_float: bool = comptime @typeInfo(From) == .float;
// VALIDATE BEFORE WRITING ANYTHING - the all-or-nothing rule `divideBy` learned the hard way,
// when a failure left `out` half-converted and the caller could not tell how far it got.
if (comptime to_int and from_float) {
const low: From = @floatFromInt(zm.minInt(To));
const high: From = @floatFromInt(zm.maxInt(To));
var check: Walk = .over(a.shape[0..a.rank]);
while (check.next()) |at| {
const x: From = a.data[a.offsetOf(at)];
if (isNan(x) or !zm.isFinite(x)) {
return Error.DomainError;
}
// `@trunc` first: 2.9 fits in an i2 whose maximum is 1 only after truncation, and
// comparing the untruncated value would reject a conversion that is exact.
const whole: From = @trunc(x);
if (whole < low or whole > high) {
return Error.DomainError;
}
}
}
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
const x: From = a.data[a.offsetOf(at)];
const converted: To = switch (@typeInfo(To)) {
.int => switch (@typeInfo(From)) {
.int => @intCast(x),
// `@trunc` and NOT `@intFromFloat`: the linter's `int-from-float` rule wants
// the rounding NAMED, because `@intFromFloat` alone leaves a reader to remember
// which of four roundings it does. `@trunc` says toward zero, and the coercion
// to `To` carries it the rest of the way.
else => @trunc(x),
},
else => switch (@typeInfo(From)) {
.int => @floatFromInt(x),
else => @floatCast(x),
},
};
out.data[out.offsetOf(at)] = converted;
}
}
Four conversions, four different hazards:
f64 -> f32 1e300 becomes infinity
f64 -> i32 3.7 becomes 3 or 4, depending on a rounding rule
f64 -> i8 300.0 is OUT OF RANGE
i32 -> f32 2^24 + 1 loses its last bit3.7 becomes 3 and -3.7 becomes -3. That is
@trunc's rule, C's and numpy's — and the negative case is where it differs from
floor, which would give −4.
A caller wanting different rounding applies it first:
round then astype gives 4. The choice ends up at the call site instead
of hidden inside a conversion.
numpy wraps silently. That is how a temperature of 300 becomes 44 in an i8 column and nobody notices for a week.
Zig cannot follow it even if it wanted to: @intFromFloat out of
range is undefined behaviour, not a wrap. So the choice was forced — and forcing it
produced the better answer. astype validates the whole tensor before writing anything,
so a refused conversion leaves the output exactly as it was.
NaN and infinity are refused for the same reason: neither has an integer, and returning whatever the hardware produces would be inventing one.
One case that is accepted: 127.9 into an i8. It does not fit as a real
number and does after truncation — which is precisely what truncation means, so rejecting it
would refuse a correct conversion.
An f32 cannot hold 1e300 and inf is the IEEE answer, not a mistake to report. The
asymmetry is deliberate: an out-of-range integer conversion has no defined result, an
out-of-range float conversion has one and it is infinity.
zn.Rng is not a stream you advance. Every draw is a pure function of a seed and an
index, so you address values rather than consume them:
const rng: zn.Rng = zn.Rng.init(7);
const a: u32 = rng.bits(900); // no draws 0..899 were madeRng/// A reproducible random source addressed by (seed, index) rather than by mutation.
///
/// WHY COUNTER-BASED, AND NOT A STATEFUL GENERATOR
///
/// A classic PRNG carries state and advances it, so drawing the i-th value means drawing the
/// first i. That is fine on one CPU thread and useless everywhere else this library is going:
///
/// * **A GPU kernel has no shared state.** Thread `i` must be able to compute ITS value without
/// any other thread having run. A counter-based generator is a pure function of the index, so
/// the kernel is `value(seed, thread_id)` and needs no sequencing at all.
/// * **CPU and GPU must agree EXACTLY.** The CPU path is this library's oracle (plan section 5 rung
/// 1); an oracle that produces a different random stream than the thing it checks is not an
/// oracle. Same function, same integers, same answer.
/// * **A failing case is reproducible from two numbers.** "seed 7, index 4096" is the whole
/// state, so a bad draw can be re-created in a unit test without replaying a sequence.
///
/// THE MIXER IS DELIBERATELY MODEST AND THE TESTS SAY SO. `mix32` is the murmur3 finalizer -
/// a bijection with good avalanche, applied to `index` folded with `seed`. It is NOT a
/// cryptographic generator and NOT claimed to pass a full statistical battery; what it is pinned
/// on is the properties this library actually relies on: determinism, range, stream independence,
/// and first/second moments. A stronger generator can replace the body without touching a single
/// call site, precisely because the interface is (seed, index) rather than a stream.
pub const Rng = struct {
seed: u32,
/// A generator for a given seed. Two `Rng`s with the same seed are the same function.
pub fn init(seed: u32) Rng {
return .{ .seed = seed };
}
/// A DERIVED generator, so one seed can feed several independent uses without the caller
/// inventing seeds. `parent.split(0)` and `parent.split(1)` are unrelated streams.
///
/// Mixing the label rather than adding it: `seed + 1` and `seed + 2` are adjacent inputs to
/// the mixer, and adjacency is exactly what a hash is asked not to preserve. Passing the
/// label through `mix32` first means the derived seeds are as far apart as any other pair.
pub fn split(self: Rng, label: u32) Rng {
return .{ .seed = mix32(self.seed ^ mix32(label ^ 0x9e3779b9)) };
}
/// The number of mantissa bits `T` can hold, including the implicit leading one. This is how
/// many random bits a draw in `[0, 1)` may use: fewer wastes range, MORE ROUNDS UP TO 1.0 and
/// breaks the half-open interval.
fn mantissaWidth(comptime T: type) u32 {
comptime requireFloat(T);
return switch (T) {
f16 => 11,
f32 => 24,
f64 => 53,
else => unreachable,
};
}
/// One 32-bit word of the generator's output at `index`, from independent stream `stream`.
///
/// EACH KIND OF DRAW OWNS A STREAM, SO KINDS DO NOT INTERFERE
///
/// Without this, every method would start from `bits(index)` and drawing two different KINDS
/// at one index would return values that determine each other - `intBelow(i, n)` would be a
/// function of `unitFloat(i)`. A caller has no reason to expect that and every reason to be
/// hurt by it, so the stream is part of the addressing rather than the caller's problem.
///
/// The stream index is folded into the MIXER, not into the index. Folding it into the index
/// (`bits(index ^ K)`) makes indices `i` and `i ^ K` draw from the same two words with their
/// roles swapped, so those pairs are not independent - which is exactly the defect this
/// replaced.
///
/// `stream == 0` adds nothing, so `bits` is unchanged by the introduction of rounds.
fn word(self: Rng, index: u32, stream: u32) u32 {
return mix32(mix32(index *% 0x9e3779b9) ^ self.seed ^ (stream *% 0x85ebca6b));
}
/// The raw 32-bit draw at `index`, stream 0. Distinct from every other kind of draw.
pub fn bits(self: Rng, index: u32) u32 {
return self.word(index, 0);
}
/// A float in `[0, 1)` from stream `stream` (and `stream + 1` when `T` needs more than 32 bits).
///
/// Exactly `mantissaWidth(T)` random bits, scaled by 2^-width. The largest result is
/// (2^width - 1) / 2^width, which is the largest `T` strictly below one - so the interval is
/// half-open for every width, including `f16`, where using 24 bits rounds to 1.0 instead.
fn unitFloatRound(self: Rng, comptime T: type, index: u32, comptime stream: u32) T {
comptime requireFloat(T);
const width: u32 = comptime mantissaWidth(T);
const first: u64 = self.word(index, stream);
const mantissa: u64 = if (comptime width <= 32)
first >> comptime @as(u6, @intCast(32 - width))
else
(first << comptime @as(u6, @intCast(width - 32))) |
(@as(u64, self.word(index, stream + 1)) >> comptime @as(u6, @intCast(64 - width)));
const numerator: T = @floatFromInt(mantissa);
const denominator: T = @floatFromInt(@as(u64, 1) << comptime @as(u6, @intCast(width)));
return numerator / denominator;
}
/// A float in `[0, 1)`, drawn at `index`.
pub fn unitFloat(self: Rng, comptime T: type, index: u32) T {
return self.unitFloatRound(T, index, 1);
}
/// A float in `[low, high)`, drawn at `index`.
pub fn uniform(
self: Rng,
comptime T: type,
index: u32,
low: T,
high: T,
) T {
comptime requireFloat(T);
return low + (high - low) * self.unitFloat(T, index);
}
/// A standard normal draw (mean 0, variance 1) at `index`, by the Box-Muller transform.
///
/// ITS TWO UNIFORMS COME FROM ITS OWN STREAMS (3-6), not from the uniform stream, so
/// `normal(i)` and `unitFloat(i)` are independent and a caller mixing them needs no care.
///
/// Box-Muller rather than a ziggurat because a ziggurat REJECTS: the number of uniforms it
/// consumes varies per sample, which would make the value at `index` depend on how many
/// rejections happened. That is not a pure function of the index and would not survive being
/// evaluated on an independent GPU thread. Box-Muller consumes exactly two, always.
///
/// The radius draw is taken from `(0, 1]` rather than `[0, 1)` because `@log(0)` is -inf.
pub fn normal(self: Rng, comptime T: type, index: u32) T {
comptime requireFloat(T);
const two_pi: T = 6.283185307179586;
const radius_draw: T = 1.0 - self.unitFloatRound(T, index, 3);
const angle_draw: T = self.unitFloatRound(T, index, 5);
const radius: T = @sqrt(-2.0 * @log(radius_draw));
return radius * @cos(two_pi * angle_draw);
}
/// An integer in `[0, bound)`, drawn at `index` from its own stream. `bound == 0` yields 0.
///
/// Lemire's multiply-shift, not `% bound`: the modulo is biased toward small values whenever
/// `bound` does not divide 2^32, and the bias stays invisible until something like a replay
/// buffer over-samples its early entries.
pub fn intBelow(self: Rng, index: u32, bound: u32) u32 {
if (bound == 0) {
return 0;
}
const wide: u64 = @as(u64, self.word(index, 7)) * @as(u64, bound);
return @intCast(wide >> 32);
}
/// Fill `out` with uniforms in `[0, 1)`, drawn at indices `0..out.len`.
pub fn fillUniform(self: Rng, comptime T: type, out: []T) void {
comptime requireFloat(T);
for (out, 0..) |*slot, i| {
slot.* = self.unitFloat(T, @intCast(i));
}
}
/// Fill `out` with standard normal draws, at indices `0..out.len`.
pub fn fillNormal(self: Rng, comptime T: type, out: []T) void {
comptime requireFloat(T);
for (out, 0..) |*slot, i| {
slot.* = self.normal(T, @intCast(i));
}
}
};
Three consequences:
value(seed, thread_id) produces the identical
stream.One 32-bit word comes from mixing the index, the seed and a stream number:
fn word(self: Rng, index: u32, stream: u32) u32 {
return mix32(mix32(index *% 0x9e3779b9) ^ self.seed ^ (stream *% 0x85ebca6b));
}mix32 is the 32-bit MurmurHash3 finaliser. Each kind of draw owns a
stream, so draws of different kinds at the same index are independent:
| Stream | Used by |
|---|---|
| 0 | bits |
| 1, 2 | unitFloat (the second word supplies the extra mantissa bits of f64) |
| 3–6 | normal, whose two uniforms each take a pair |
| 7 | intBelow |
The stream is folded into the mixer rather than into the index. Folding it into the index
would make indices i and i ^ K draw from the same words with their roles
swapped.
unitFloat takes exactly as many random bits as T's mantissa holds
— 11 for f16, 24 for f32, 53 for f64 — and
scales by 2−width. The largest result is the largest T strictly
below one, so the interval is half-open at every width.
const rng: zn.Rng = zn.Rng.init(20260904);
const raw: u32 = rng.bits(0); // raw 32 bits
const u: f64 = rng.unitFloat(f64, 0); // [0, 1)
const t: f32 = rng.uniform(f32, 0, -1.0, 1.0); // [-1, 1)
const z: f64 = rng.normal(f64, 0); // standard normal
const k: u32 = rng.intBelow(0, 52); // [0, 52), unbiased
var buf: [256]f32 = undefined;
rng.fillUniform(f32, &buf); // indices 0..255
rng.fillNormal(f32, &buf); // indices 0..255Rng/// A reproducible random source addressed by (seed, index) rather than by mutation.
///
/// WHY COUNTER-BASED, AND NOT A STATEFUL GENERATOR
///
/// A classic PRNG carries state and advances it, so drawing the i-th value means drawing the
/// first i. That is fine on one CPU thread and useless everywhere else this library is going:
///
/// * **A GPU kernel has no shared state.** Thread `i` must be able to compute ITS value without
/// any other thread having run. A counter-based generator is a pure function of the index, so
/// the kernel is `value(seed, thread_id)` and needs no sequencing at all.
/// * **CPU and GPU must agree EXACTLY.** The CPU path is this library's oracle (plan section 5 rung
/// 1); an oracle that produces a different random stream than the thing it checks is not an
/// oracle. Same function, same integers, same answer.
/// * **A failing case is reproducible from two numbers.** "seed 7, index 4096" is the whole
/// state, so a bad draw can be re-created in a unit test without replaying a sequence.
///
/// THE MIXER IS DELIBERATELY MODEST AND THE TESTS SAY SO. `mix32` is the murmur3 finalizer -
/// a bijection with good avalanche, applied to `index` folded with `seed`. It is NOT a
/// cryptographic generator and NOT claimed to pass a full statistical battery; what it is pinned
/// on is the properties this library actually relies on: determinism, range, stream independence,
/// and first/second moments. A stronger generator can replace the body without touching a single
/// call site, precisely because the interface is (seed, index) rather than a stream.
pub const Rng = struct {
seed: u32,
/// A generator for a given seed. Two `Rng`s with the same seed are the same function.
pub fn init(seed: u32) Rng {
return .{ .seed = seed };
}
/// A DERIVED generator, so one seed can feed several independent uses without the caller
/// inventing seeds. `parent.split(0)` and `parent.split(1)` are unrelated streams.
///
/// Mixing the label rather than adding it: `seed + 1` and `seed + 2` are adjacent inputs to
/// the mixer, and adjacency is exactly what a hash is asked not to preserve. Passing the
/// label through `mix32` first means the derived seeds are as far apart as any other pair.
pub fn split(self: Rng, label: u32) Rng {
return .{ .seed = mix32(self.seed ^ mix32(label ^ 0x9e3779b9)) };
}
/// The number of mantissa bits `T` can hold, including the implicit leading one. This is how
/// many random bits a draw in `[0, 1)` may use: fewer wastes range, MORE ROUNDS UP TO 1.0 and
/// breaks the half-open interval.
fn mantissaWidth(comptime T: type) u32 {
comptime requireFloat(T);
return switch (T) {
f16 => 11,
f32 => 24,
f64 => 53,
else => unreachable,
};
}
/// One 32-bit word of the generator's output at `index`, from independent stream `stream`.
///
/// EACH KIND OF DRAW OWNS A STREAM, SO KINDS DO NOT INTERFERE
///
/// Without this, every method would start from `bits(index)` and drawing two different KINDS
/// at one index would return values that determine each other - `intBelow(i, n)` would be a
/// function of `unitFloat(i)`. A caller has no reason to expect that and every reason to be
/// hurt by it, so the stream is part of the addressing rather than the caller's problem.
///
/// The stream index is folded into the MIXER, not into the index. Folding it into the index
/// (`bits(index ^ K)`) makes indices `i` and `i ^ K` draw from the same two words with their
/// roles swapped, so those pairs are not independent - which is exactly the defect this
/// replaced.
///
/// `stream == 0` adds nothing, so `bits` is unchanged by the introduction of rounds.
fn word(self: Rng, index: u32, stream: u32) u32 {
return mix32(mix32(index *% 0x9e3779b9) ^ self.seed ^ (stream *% 0x85ebca6b));
}
/// The raw 32-bit draw at `index`, stream 0. Distinct from every other kind of draw.
pub fn bits(self: Rng, index: u32) u32 {
return self.word(index, 0);
}
/// A float in `[0, 1)` from stream `stream` (and `stream + 1` when `T` needs more than 32 bits).
///
/// Exactly `mantissaWidth(T)` random bits, scaled by 2^-width. The largest result is
/// (2^width - 1) / 2^width, which is the largest `T` strictly below one - so the interval is
/// half-open for every width, including `f16`, where using 24 bits rounds to 1.0 instead.
fn unitFloatRound(self: Rng, comptime T: type, index: u32, comptime stream: u32) T {
comptime requireFloat(T);
const width: u32 = comptime mantissaWidth(T);
const first: u64 = self.word(index, stream);
const mantissa: u64 = if (comptime width <= 32)
first >> comptime @as(u6, @intCast(32 - width))
else
(first << comptime @as(u6, @intCast(width - 32))) |
(@as(u64, self.word(index, stream + 1)) >> comptime @as(u6, @intCast(64 - width)));
const numerator: T = @floatFromInt(mantissa);
const denominator: T = @floatFromInt(@as(u64, 1) << comptime @as(u6, @intCast(width)));
return numerator / denominator;
}
/// A float in `[0, 1)`, drawn at `index`.
pub fn unitFloat(self: Rng, comptime T: type, index: u32) T {
return self.unitFloatRound(T, index, 1);
}
/// A float in `[low, high)`, drawn at `index`.
pub fn uniform(
self: Rng,
comptime T: type,
index: u32,
low: T,
high: T,
) T {
comptime requireFloat(T);
return low + (high - low) * self.unitFloat(T, index);
}
/// A standard normal draw (mean 0, variance 1) at `index`, by the Box-Muller transform.
///
/// ITS TWO UNIFORMS COME FROM ITS OWN STREAMS (3-6), not from the uniform stream, so
/// `normal(i)` and `unitFloat(i)` are independent and a caller mixing them needs no care.
///
/// Box-Muller rather than a ziggurat because a ziggurat REJECTS: the number of uniforms it
/// consumes varies per sample, which would make the value at `index` depend on how many
/// rejections happened. That is not a pure function of the index and would not survive being
/// evaluated on an independent GPU thread. Box-Muller consumes exactly two, always.
///
/// The radius draw is taken from `(0, 1]` rather than `[0, 1)` because `@log(0)` is -inf.
pub fn normal(self: Rng, comptime T: type, index: u32) T {
comptime requireFloat(T);
const two_pi: T = 6.283185307179586;
const radius_draw: T = 1.0 - self.unitFloatRound(T, index, 3);
const angle_draw: T = self.unitFloatRound(T, index, 5);
const radius: T = @sqrt(-2.0 * @log(radius_draw));
return radius * @cos(two_pi * angle_draw);
}
/// An integer in `[0, bound)`, drawn at `index` from its own stream. `bound == 0` yields 0.
///
/// Lemire's multiply-shift, not `% bound`: the modulo is biased toward small values whenever
/// `bound` does not divide 2^32, and the bias stays invisible until something like a replay
/// buffer over-samples its early entries.
pub fn intBelow(self: Rng, index: u32, bound: u32) u32 {
if (bound == 0) {
return 0;
}
const wide: u64 = @as(u64, self.word(index, 7)) * @as(u64, bound);
return @intCast(wide >> 32);
}
/// Fill `out` with uniforms in `[0, 1)`, drawn at indices `0..out.len`.
pub fn fillUniform(self: Rng, comptime T: type, out: []T) void {
comptime requireFloat(T);
for (out, 0..) |*slot, i| {
slot.* = self.unitFloat(T, @intCast(i));
}
}
/// Fill `out` with standard normal draws, at indices `0..out.len`.
pub fn fillNormal(self: Rng, comptime T: type, out: []T) void {
comptime requireFloat(T);
for (out, 0..) |*slot, i| {
slot.* = self.normal(T, @intCast(i));
}
}
};
The method is unitFloat rather than float because
float is reserved in this tree — zm.float is the
integer-to-f32 conversion. intBelow uses Lemire's multiply-shift, not a
remainder, so it carries no bias toward small values.
The bulk fills write indices 0..out.len and are exactly equivalent to calling the
single-value method at each index, so the two forms mix freely.
const root: zn.Rng = zn.Rng.init(1);
const weights: zn.Rng = root.split(0);
const dropout: zn.Rng = root.split(1);Rng/// A reproducible random source addressed by (seed, index) rather than by mutation.
///
/// WHY COUNTER-BASED, AND NOT A STATEFUL GENERATOR
///
/// A classic PRNG carries state and advances it, so drawing the i-th value means drawing the
/// first i. That is fine on one CPU thread and useless everywhere else this library is going:
///
/// * **A GPU kernel has no shared state.** Thread `i` must be able to compute ITS value without
/// any other thread having run. A counter-based generator is a pure function of the index, so
/// the kernel is `value(seed, thread_id)` and needs no sequencing at all.
/// * **CPU and GPU must agree EXACTLY.** The CPU path is this library's oracle (plan section 5 rung
/// 1); an oracle that produces a different random stream than the thing it checks is not an
/// oracle. Same function, same integers, same answer.
/// * **A failing case is reproducible from two numbers.** "seed 7, index 4096" is the whole
/// state, so a bad draw can be re-created in a unit test without replaying a sequence.
///
/// THE MIXER IS DELIBERATELY MODEST AND THE TESTS SAY SO. `mix32` is the murmur3 finalizer -
/// a bijection with good avalanche, applied to `index` folded with `seed`. It is NOT a
/// cryptographic generator and NOT claimed to pass a full statistical battery; what it is pinned
/// on is the properties this library actually relies on: determinism, range, stream independence,
/// and first/second moments. A stronger generator can replace the body without touching a single
/// call site, precisely because the interface is (seed, index) rather than a stream.
pub const Rng = struct {
seed: u32,
/// A generator for a given seed. Two `Rng`s with the same seed are the same function.
pub fn init(seed: u32) Rng {
return .{ .seed = seed };
}
/// A DERIVED generator, so one seed can feed several independent uses without the caller
/// inventing seeds. `parent.split(0)` and `parent.split(1)` are unrelated streams.
///
/// Mixing the label rather than adding it: `seed + 1` and `seed + 2` are adjacent inputs to
/// the mixer, and adjacency is exactly what a hash is asked not to preserve. Passing the
/// label through `mix32` first means the derived seeds are as far apart as any other pair.
pub fn split(self: Rng, label: u32) Rng {
return .{ .seed = mix32(self.seed ^ mix32(label ^ 0x9e3779b9)) };
}
/// The number of mantissa bits `T` can hold, including the implicit leading one. This is how
/// many random bits a draw in `[0, 1)` may use: fewer wastes range, MORE ROUNDS UP TO 1.0 and
/// breaks the half-open interval.
fn mantissaWidth(comptime T: type) u32 {
comptime requireFloat(T);
return switch (T) {
f16 => 11,
f32 => 24,
f64 => 53,
else => unreachable,
};
}
/// One 32-bit word of the generator's output at `index`, from independent stream `stream`.
///
/// EACH KIND OF DRAW OWNS A STREAM, SO KINDS DO NOT INTERFERE
///
/// Without this, every method would start from `bits(index)` and drawing two different KINDS
/// at one index would return values that determine each other - `intBelow(i, n)` would be a
/// function of `unitFloat(i)`. A caller has no reason to expect that and every reason to be
/// hurt by it, so the stream is part of the addressing rather than the caller's problem.
///
/// The stream index is folded into the MIXER, not into the index. Folding it into the index
/// (`bits(index ^ K)`) makes indices `i` and `i ^ K` draw from the same two words with their
/// roles swapped, so those pairs are not independent - which is exactly the defect this
/// replaced.
///
/// `stream == 0` adds nothing, so `bits` is unchanged by the introduction of rounds.
fn word(self: Rng, index: u32, stream: u32) u32 {
return mix32(mix32(index *% 0x9e3779b9) ^ self.seed ^ (stream *% 0x85ebca6b));
}
/// The raw 32-bit draw at `index`, stream 0. Distinct from every other kind of draw.
pub fn bits(self: Rng, index: u32) u32 {
return self.word(index, 0);
}
/// A float in `[0, 1)` from stream `stream` (and `stream + 1` when `T` needs more than 32 bits).
///
/// Exactly `mantissaWidth(T)` random bits, scaled by 2^-width. The largest result is
/// (2^width - 1) / 2^width, which is the largest `T` strictly below one - so the interval is
/// half-open for every width, including `f16`, where using 24 bits rounds to 1.0 instead.
fn unitFloatRound(self: Rng, comptime T: type, index: u32, comptime stream: u32) T {
comptime requireFloat(T);
const width: u32 = comptime mantissaWidth(T);
const first: u64 = self.word(index, stream);
const mantissa: u64 = if (comptime width <= 32)
first >> comptime @as(u6, @intCast(32 - width))
else
(first << comptime @as(u6, @intCast(width - 32))) |
(@as(u64, self.word(index, stream + 1)) >> comptime @as(u6, @intCast(64 - width)));
const numerator: T = @floatFromInt(mantissa);
const denominator: T = @floatFromInt(@as(u64, 1) << comptime @as(u6, @intCast(width)));
return numerator / denominator;
}
/// A float in `[0, 1)`, drawn at `index`.
pub fn unitFloat(self: Rng, comptime T: type, index: u32) T {
return self.unitFloatRound(T, index, 1);
}
/// A float in `[low, high)`, drawn at `index`.
pub fn uniform(
self: Rng,
comptime T: type,
index: u32,
low: T,
high: T,
) T {
comptime requireFloat(T);
return low + (high - low) * self.unitFloat(T, index);
}
/// A standard normal draw (mean 0, variance 1) at `index`, by the Box-Muller transform.
///
/// ITS TWO UNIFORMS COME FROM ITS OWN STREAMS (3-6), not from the uniform stream, so
/// `normal(i)` and `unitFloat(i)` are independent and a caller mixing them needs no care.
///
/// Box-Muller rather than a ziggurat because a ziggurat REJECTS: the number of uniforms it
/// consumes varies per sample, which would make the value at `index` depend on how many
/// rejections happened. That is not a pure function of the index and would not survive being
/// evaluated on an independent GPU thread. Box-Muller consumes exactly two, always.
///
/// The radius draw is taken from `(0, 1]` rather than `[0, 1)` because `@log(0)` is -inf.
pub fn normal(self: Rng, comptime T: type, index: u32) T {
comptime requireFloat(T);
const two_pi: T = 6.283185307179586;
const radius_draw: T = 1.0 - self.unitFloatRound(T, index, 3);
const angle_draw: T = self.unitFloatRound(T, index, 5);
const radius: T = @sqrt(-2.0 * @log(radius_draw));
return radius * @cos(two_pi * angle_draw);
}
/// An integer in `[0, bound)`, drawn at `index` from its own stream. `bound == 0` yields 0.
///
/// Lemire's multiply-shift, not `% bound`: the modulo is biased toward small values whenever
/// `bound` does not divide 2^32, and the bias stays invisible until something like a replay
/// buffer over-samples its early entries.
pub fn intBelow(self: Rng, index: u32, bound: u32) u32 {
if (bound == 0) {
return 0;
}
const wide: u64 = @as(u64, self.word(index, 7)) * @as(u64, bound);
return @intCast(wide >> 32);
}
/// Fill `out` with uniforms in `[0, 1)`, drawn at indices `0..out.len`.
pub fn fillUniform(self: Rng, comptime T: type, out: []T) void {
comptime requireFloat(T);
for (out, 0..) |*slot, i| {
slot.* = self.unitFloat(T, @intCast(i));
}
}
/// Fill `out` with standard normal draws, at indices `0..out.len`.
pub fn fillNormal(self: Rng, comptime T: type, out: []T) void {
comptime requireFloat(T);
for (out, 0..) |*slot, i| {
slot.* = self.normal(T, @intCast(i));
}
}
};
split derives a generator from a label, so one seed feeds several uses without
inventing seeds by hand. Labels are hashed rather than added, so split(0) and
split(1) are as unrelated as any other pair.
| Guarantee | Checked by |
|---|---|
| The same (seed, index) gives the same value, in any order | Drawing one index first from one generator and last from another |
unitFloat is in [0, 1) at every float width; uniform in [low, high) | 50,000 draws each at f16, f32, f64 |
| Draws of different kinds at one index are independent | 20,000 indices, no coincidences between bits, unitFloat, intBelow and normal |
| Uniform mean and variance are 1/2 and 1/12 | 100,000 draws, tolerance about four standard errors |
| Normal mean and variance are 0 and 1 | 100,000 draws, same tolerance |
Different seeds and different split labels give different streams | Zero identical draws across 10,000 indices |
intBelow is unbiased | A bound that does not divide 232, flat to five standard deviations |
| The mixer is injective and avalanches | No collisions over 50,000 values; one input bit flips 16 ± 0.5 output bits |
The moment tolerances are floors: a healthy generator effectively never trips them and a broken one always does. Passing them says the generator is not obviously broken, not that it is good. It is not cryptographic.
const ctx: zn.Ctx = zn.Ctx.init(gpa, 4242); // { gpa, rng }
const dropout: zn.Ctx = ctx.derive(0); // same allocator, own stream
const weights: zn.Ctx = ctx.derive(1);Ctx/// What a zimrnum call needs from its caller: somewhere to allocate, and a source of randomness.
///
/// WHY THIS IS A PARAMETER AND NOT A GLOBAL
///
/// A numerical library has exactly two ambient resources it is tempted to hide, and hiding either
/// one costs more than it saves. A module-level allocator makes every function's memory behaviour
/// invisible at the call site and untestable under a leak-checking allocator. A module-level
/// random state makes results depend on evaluation order, which is the property this library's
/// generator was specifically designed not to have (see `Rng`).
///
/// zimr forbids the first outright - the `module-var` lint rejects mutable file-scope state - and
/// `Ctx` is how the second stays honest too. A function that takes a `Ctx` says what it uses.
pub const Ctx = struct {
gpa: Allocator,
rng: Rng,
/// A context over `gpa`, with a generator seeded by `seed`.
pub fn init(gpa: Allocator, seed: u32) Ctx {
return .{ .gpa = gpa, .rng = Rng.init(seed) };
}
/// A context sharing this allocator, with a DERIVED generator. Use it when a sub-computation
/// needs its own random stream and must not consume this one's - `Rng.split` is what makes
/// that free of coordination.
pub fn derive(self: Ctx, label: u32) Ctx {
return .{ .gpa = self.gpa, .rng = self.rng.split(label) };
}
};
A Ctx carries the two things a numerical call needs from its caller: somewhere to
allocate and a source of randomness. Both are parameters rather than globals, so a function's
memory behaviour is visible at its call site and testable under a leak-checking allocator, and its
results do not depend on what else has drawn.
var arena: std.heap.ArenaAllocator = .init(gpa);
defer arena.deinit();
const a: Allocator = arena.allocator();
const buffer: []f64 = try a.alloc(f64, 4096);
rng.fillNormal(f64, buffer);
// no free for `buffer`: deinit releases the arena and everything in itAn arena covers one unit of work — a training step, one inference, one solve —
released as a whole. Intermediates inside such a unit have exactly the arena's lifetime, so they
need no individual defer.
The caller owns it. zimrnum used to wrap this in a Scope holding
*std.heap.ArenaAllocator, heap-allocated so the Allocator interface stayed
valid when the Scope value moved. It was removed: three lines either way, and one
fewer heap allocation.
The hazard the wrapper guarded against is real but narrow, and worth knowing. Measured:
by-value arena moved, interface RE-MINTED passes
by-value arena moved, interface CACHED ABORTS
Do not cache an Allocator across a move of the value that made it.
Re-mint from the arena where it now lives. No caller in this library ever cached one, which is why
the wrapper was paying for a hazard nothing had.
A struct that must own an arena takes an out pointer:
fn initState(gpa: Allocator, s: *State) !void {
s.arena = .init(gpa); // written to its FINAL address
const a = s.arena.allocator(); // minted from there
...
}Nothing moves after the interface exists, so this is safer than a movable wrapper rather than merely simpler. It is also what the Zig standard library expects of you.
var arena: std.heap.ArenaAllocator = .init(gpa);
defer arena.deinit();
for (0..epochs) |_| {
trainOneStep(arena.allocator());
_ = arena.reset(.retain_capacity); // frees the step, keeps the memory
}reset(.retain_capacity) releases the contents and retains the memory. For a loop
whose iterations do the same shape of work, the first grows the arena to its high-water mark and
every one after reuses it — so a training loop allocates from the parent once.
Note that arena.allocator() is called inside the loop. That costs nothing
and sidesteps the caching hazard above entirely.
// From an arena: nothing to free.
var arena: std.heap.ArenaAllocator = .init(gpa);
defer arena.deinit();
const t = try zn.Tensor(f64).alloc(s.allocator(), &.{ 8, 8 });
// From an allocator you manage:
const u = try zn.Tensor(f32).alloc(gpa, &.{ 2, 3, 4 });
defer gpa.free(u.data);
// Over storage you already have:
var backing: [10]f32 = @splat(0);
const v = try zn.Tensor(f32).fromSlice(&backing, &.{ 2, 5 });Tensor/// A view of `T`-typed storage as a multidimensional array.
///
/// A TENSOR NEVER OWNS ITS STORAGE, AND THAT IS THE DESIGN
///
/// There is no ownership flag and no `deinit`. `data` is a slice the tensor addresses; whoever
/// allocated it frees it, with the allocator that allocated it. Usually that is an arena, in
/// which case nobody frees anything and `arena.deinit()` ends the matter.
///
/// The alternative - an `owns: bool` and a `deinit` that is a no-op for views - is the shape
/// this library was warned about before it had a tensor: both of its failure modes are SILENT.
/// `deinit` on a view does nothing and reads as if it did something; a missing `deinit` on an
/// owner just grows memory. Having no `deinit` at all means a free appears exactly where an
/// allocation appears, and a view has neither.
///
/// Strides are in ELEMENTS and signed, so a reversed or negatively-strided view is expressible
/// without changing this type. `at` computes a signed offset and casts once.
pub fn Tensor(comptime T: type) type {
return struct {
const Self = @This();
/// The storage this tensor addresses. Not owned - see the type's doc comment.
data: []T,
/// Extent of each axis. Only the first `rank` entries are meaningful.
shape: [max_rank]usize,
/// Step in ELEMENTS between consecutive positions along each axis.
strides: [max_rank]isize,
rank: u8,
/// WHERE THIS VIEW STARTS INSIDE `data`
///
/// Every view until now began at element 0, which is why `slice` could not exist: a
/// window into the middle of a tensor IS an offset. With this, `slice` and the rest of
/// znum's `index` namespace become views rather than copies.
///
/// IT ALSO WEAKENS WHAT `data.len` MEANS. `data` is the whole backing store; `base`
/// and the strides say which of it this view addresses. Everything that walked
/// `data[0..size()]` assuming those coincide had to be corrected - five sites, all of
/// them fast paths that skip the index walk, which is exactly where such an assumption
/// hides.
base: usize = 0,
/// A dense row-major tensor over `shape`, allocated from `gpa`.
///
/// The caller frees `result.data` with the same allocator - or allocates from an arena
/// and frees nothing.
pub fn alloc(gpa: Allocator, shape: []const usize) Error!Self {
if (shape.len > max_rank) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
const storage: []T = gpa.alloc(T, count) catch return Error.OutOfMemory;
return fromSlice(storage, shape);
}
/// A dense row-major view of `data` as `shape`. `data.len` must be exactly the product of
/// the extents - a slice that is merely LARGE ENOUGH is rejected, because the usual reason
/// for the mismatch is a shape that is wrong rather than storage that is generous.
pub fn fromSlice(data: []T, shape: []const usize) Error!Self {
if (shape.len > max_rank) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
if (count != data.len) {
return Error.ShapeMismatch;
}
var self: Self = .{
.data = data,
.shape = @splat(1),
.strides = @splat(0),
.rank = @intCast(shape.len),
.base = 0,
};
// Row-major: the last axis is contiguous, and each earlier stride is the product of
// the extents to its right.
var stride: isize = 1;
var axis: usize = shape.len;
while (axis > 0) {
axis -= 1;
self.shape[axis] = shape[axis];
self.strides[axis] = stride;
stride *= @intCast(shape[axis]);
}
return self;
}
/// The number of elements this tensor addresses.
pub fn size(self: Self) usize {
var count: usize = 1;
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
count *= self.shape[axis];
}
return count;
}
/// Whether the strides are dense row-major, so `data[0..size()]` is the tensor in order.
///
/// An operation that walks memory linearly must ask this rather than assume it: a view
/// produced by `transpose` addresses the same storage in a different order, and a linear
/// walk over it silently reads the wrong elements rather than failing.
pub fn isContiguous(self: Self) bool {
var expected: isize = 1;
var axis: usize = self.rank;
while (axis > 0) {
axis -= 1;
if (self.shape[axis] != 1 and self.strides[axis] != expected) {
return false;
}
expected *= @intCast(self.shape[axis]);
}
return true;
}
/// The offset into `data` of the element at `indices`.
pub fn flatIndex(self: Self, indices: []const usize) Error!usize {
if (indices.len != self.rank) {
return Error.ShapeMismatch;
}
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
if (i >= self.shape[axis]) {
return Error.OutOfRange;
}
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}
/// The element at `indices`.
pub fn at(self: Self, indices: []const usize) Error!T {
return self.data[try self.flatIndex(indices)];
}
/// The flat offset for a coordinate, with no error.
///
/// THE 45 `flatIndex(...) catch unreachable` SITES ARE ALL THIS SHAPE
///
/// const at: []const usize = walker[0..out.rank];
/// out.data[out.offsetOf(at)] = ...
///
/// The slice is `walker[0..rank]`, so `indices.len != rank` **cannot** happen - the
/// caller built it from the rank. And the coordinate came from a `Walk` over this very
/// shape, so the bounds cannot fail either. Two impossible checks, one `catch
/// unreachable`, on the hottest line in the library.
///
/// This computes the offset and returns it. The remaining safety is `data[offset]`,
/// which is where Zig puts it: a panic with a message in Debug and ReleaseSafe.
///
/// **Use it where the rank is an invariant you have already established** - after a
/// shape check, or inside a walk over the tensor's own shape. Use `flatIndex` where the
/// coordinate came from outside and might be wrong.
pub fn offsetOf(self: Self, indices: []const usize) usize {
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}
/// Read from a rank-2 tensor. **Does not return an error.**
///
/// WHY `at` HAS TO AND THIS DOES NOT
///
/// `at(indices: []const usize)` can fail for two reasons. One is a real bounds check.
/// The other is `indices.len != rank` - and that check exists ONLY because a slice hides
/// its arity from the type system. Every call site writes a literal `&.{ i, j }`; the
/// compiler could know it is two, and does not.
///
/// Naming the arity in the function removes that half entirely. What remains is the
/// bounds check, and **Zig already has a name for that**: it is what `data[idx]` does -
/// a panic with a message and a stack trace in Debug and ReleaseSafe, the same UB as any
/// other slice index in ReleaseFast.
///
/// Nobody writes `arr[i] catch unreachable`. This file had **105 `catch unreachable`**,
/// most of them on an `at` whose rank the author had checked three lines earlier, and
/// that noise is the single biggest thing that makes numeric Zig look worse than it is.
///
/// The safety is not weaker - it is Zig's, in the place Zig puts it.
pub fn at2(self: Self, i: usize, j: usize) T {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i)) +
self.strides[1] * @as(isize, @intCast(j));
return self.data[@intCast(offset)];
}
/// Write to a rank-2 tensor. The mirror of `at2`, and error-free for the same reason.
pub fn setAt2(self: Self, i: usize, j: usize, value: T) void {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i)) +
self.strides[1] * @as(isize, @intCast(j));
self.data[@intCast(offset)] = value;
}
/// Read from a rank-1 tensor, error-free.
pub fn at1(self: Self, i: usize) T {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i));
return self.data[@intCast(offset)];
}
/// Write to a rank-1 tensor, error-free.
pub fn setAt1(self: Self, i: usize, value: T) void {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i));
self.data[@intCast(offset)] = value;
}
/// Write `value` at `indices`.
pub fn setAt(self: Self, indices: []const usize, value: T) Error!void {
self.data[try self.flatIndex(indices)] = value;
}
/// Set every element to `value`.
pub fn fill(self: Self, value: T) void {
if (self.isContiguous()) {
@memset(self.data[self.base..][0..self.size()], value);
return;
}
var walker: [max_rank]usize = @splat(0);
var remaining: usize = self.size();
while (remaining > 0) : (remaining -= 1) {
self.data[self.offsetOf(walker[0..self.rank])] = value;
advance(self.shape[0..self.rank], walker[0..self.rank]);
}
}
/// The same storage seen as `shape`. Requires a contiguous tensor and an equal element
/// count; both are checked rather than assumed.
pub fn reshape(self: Self, shape: []const usize) Error!Self {
if (!self.isContiguous()) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
if (count != self.size()) {
return Error.ShapeMismatch;
}
// From `base`, not 0: reshaping a slice must keep addressing that slice.
return fromSlice(self.data[self.base..][0..self.size()], shape);
}
/// Whether two distinct index positions address the same element, which happens exactly
/// when a stretched axis (stride 0) has extent above 1.
///
/// An operation that WRITES must refuse an aliased destination. Many indices landing on
/// one element means the last write wins and every earlier one vanishes - a result that is
/// wrong in a way no bounds check catches. `map` and `zip` check this so a caller who
/// passes a broadcast view as an output gets an error rather than a plausible answer.
pub fn isAliased(self: Self) bool {
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
if (self.shape[axis] > 1 and self.strides[axis] == 0) {
return true;
}
}
return false;
}
/// A read-only view of this tensor stretched to `shape` by the broadcasting rule.
///
/// BROADCASTING IS A STRIDE OF ZERO, NOT A COPY
///
/// An axis whose extent is 1 is stretched by setting its stride to 0, so every position
/// along it addresses the same element. No storage is allocated and no data moves; a
/// `(3, 1, 5)` tensor viewed as `(3, 4, 5)` is the same 15 values read four times.
///
/// **WRITING THROUGH A STRETCHED VIEW IS ALMOST CERTAINLY A MISTAKE.** Many indices
/// alias one element, so the last write wins and the others vanish. This returns a view
/// because that is what the type system has; treat it as an input.
///
/// `isContiguous` reports false for any stretched axis with extent > 1, which is what
/// stops a linear walk from reading the underlying 15 values as if they were 60.
pub fn broadcastTo(self: Self, shape: []const usize) Error!Self {
if (shape.len > max_rank or shape.len < self.rank) {
return Error.ShapeMismatch;
}
var view: Self = .{
.data = self.data,
.shape = @splat(1),
.strides = @splat(0),
.rank = @intCast(shape.len),
// Stretching does not move where the view starts.
.base = self.base,
};
// Align at the TRAILING axis, as the broadcast rule does: an axis this tensor does not
// have behaves as extent 1, and so does an axis it has with extent 1. Both get a
// stride of 0.
for (shape, 0..) |extent, axis| {
const from_end: usize = shape.len - axis;
view.shape[axis] = extent;
if (self.rank < from_end) {
continue;
}
const own_axis: usize = self.rank - from_end;
const own_extent: usize = self.shape[own_axis];
if (own_extent == extent) {
view.strides[axis] = self.strides[own_axis];
} else if (own_extent != 1) {
return Error.ShapeMismatch;
}
}
return view;
}
/// A window into this tensor: `len` positions along `axis`, starting at `start`.
///
/// A SLICE IS AN OFFSET, WHICH IS WHY `base` EXISTS
///
/// Nothing is copied: the view keeps every stride and simply starts `start` steps further
/// into the data along one axis. A batch of 32 rows out of 4096 is a `Tensor` the same
/// size as any other, addressing 32 rows of the same buffer.
///
/// `start + len` must fit. Clamping instead would hand back a view SMALLER than asked
/// for, and the caller's next loop would read the shape it requested rather than the one
/// it got.
pub fn slice(self: Self, axis: usize, start: usize, len: usize) Error!Self {
if (axis >= self.rank) {
return Error.OutOfRange;
}
if (start + len > self.shape[axis]) {
return Error.OutOfRange;
}
var view: Self = self;
const step: isize = self.strides[axis] * @as(isize, @intCast(start));
const moved: isize = @as(isize, @intCast(self.base)) + step;
if (moved < 0) {
return Error.OutOfRange;
}
view.base = @intCast(moved);
view.shape[axis] = len;
return view;
}
/// A view with the axes reordered: `out` axis `i` is this tensor's axis `order[i]`.
///
/// EVERY `manip` OPERATION HERE IS METADATA
///
/// `permute`, `squeeze`, `unsqueeze`, `flatten` and `moveAxis` all return a view: they
/// permute or insert entries in two fixed-size arrays and copy nothing. That is why they
/// need no GPU kernel and no sweep row - **there is no arithmetic to disagree about.**
/// The cost appears later, when a permuted view has to be BOUND to a kernel and must be
/// densified first, which is the constraint recorded in section 10.3.
///
/// `order` must be a permutation of `0..rank`. A repeated axis would produce a view
/// where two indices alias, which is `isAliased`'s job to reject at write time - but
/// catching it here names the actual mistake.
pub fn permute(self: Self, order: []const usize) Error!Self {
if (order.len != self.rank) {
return Error.ShapeMismatch;
}
var seen: [max_rank]bool = @splat(false);
for (order) |axis| {
if (axis >= self.rank or seen[axis]) {
return Error.OutOfRange;
}
seen[axis] = true;
}
var view: Self = self;
for (order, 0..) |axis, i| {
view.shape[i] = self.shape[axis];
view.strides[i] = self.strides[axis];
}
return view;
}
/// A view with a size-1 axis removed.
///
/// Removing an axis whose extent is NOT 1 would change how many elements the view
/// addresses, so it is `ShapeMismatch` rather than a silent reshape.
pub fn squeeze(self: Self, axis: usize) Error!Self {
if (axis >= self.rank) {
return Error.OutOfRange;
}
if (self.shape[axis] != 1) {
return Error.ShapeMismatch;
}
var view: Self = self;
var i: usize = axis;
while (i + 1 < self.rank) : (i += 1) {
view.shape[i] = self.shape[i + 1];
view.strides[i] = self.strides[i + 1];
}
view.shape[self.rank - 1] = 1;
view.strides[self.rank - 1] = 0;
view.rank = @intCast(self.rank - 1);
return view;
}
/// A view with a size-1 axis inserted at `axis`. `axis == rank` appends.
///
/// The new axis gets a stride of 0, which is harmless because its extent is 1 - a
/// single position cannot step anywhere. `isAliased` only objects when extent exceeds 1.
pub fn unsqueeze(self: Self, axis: usize) Error!Self {
if (axis > self.rank or self.rank + 1 > max_rank) {
return Error.OutOfRange;
}
var view: Self = self;
var i: usize = self.rank;
while (i > axis) : (i -= 1) {
view.shape[i] = self.shape[i - 1];
view.strides[i] = self.strides[i - 1];
}
view.shape[axis] = 1;
view.strides[axis] = 0;
view.rank = @intCast(self.rank + 1);
return view;
}
/// A rank-1 view of the whole tensor. Contiguous only, for `reshape`'s reason.
pub fn flatten(self: Self) Error!Self {
return self.reshape(&.{self.size()});
}
/// A view with one axis moved from `from` to `to`, the others keeping their order.
pub fn moveAxis(self: Self, from: usize, to: usize) Error!Self {
if (from >= self.rank or to >= self.rank) {
return Error.OutOfRange;
}
var order: [max_rank]usize = @splat(0);
var k: usize = 0;
var i: usize = 0;
while (i < self.rank) : (i += 1) {
if (i == from) {
continue;
}
order[k] = i;
k += 1;
}
// Open a slot at `to` and drop the moved axis into it.
var j: usize = self.rank - 1;
while (j > to) : (j -= 1) {
order[j] = order[j - 1];
}
order[to] = from;
return self.permute(order[0..self.rank]);
}
/// A view with axes `a` and `b` exchanged. Shares storage; changes through one are visible
/// through the other.
pub fn transpose(self: Self, a: usize, b: usize) Error!Self {
if (a >= self.rank or b >= self.rank) {
return Error.OutOfRange;
}
var view: Self = self;
view.shape[a] = self.shape[b];
view.shape[b] = self.shape[a];
view.strides[a] = self.strides[b];
view.strides[b] = self.strides[a];
return view;
}
};
}
A tensor never owns its storage. There is no ownership flag and no
deinit. t.data is a slice the tensor addresses; whoever allocated it
frees it, and if that was a Scope, nobody frees anything. A free therefore appears
exactly where an allocation appears, and a view has neither.
fromSlice requires the extents to multiply to exactly data.len. A
slice that is merely large enough is rejected.
A tensor is four public fields and no allocation of its own:
data: []T, // the backing store, not owned
shape: [max_rank]usize, // extent of each axis
strides: [max_rank]isize, // step in ELEMENTS along each axis
rank: u8,
base: usize, // where this view starts inside `data`Every read and write goes through one function:
pub fn flatIndex(self: Self, indices: []const usize) Error!usize {
if (indices.len != self.rank) return Error.ShapeMismatch;
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
if (i >= self.shape[axis]) return Error.OutOfRange;
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}That is the whole addressing model, and it explains every view operation in this section.
alloc and fromSlice produce row-major strides — the last axis steps by
1, each earlier stride is the product of the extents to its right — with base zero.
Everything else permutes those numbers.
Strides are signed and in elements. max_rank is 6, so shape and strides are
fixed-size arrays and copying a Tensor is free.
try t.setAt(&.{ 1, 2 }, 3.5);
const x: f64 = try t.at(&.{ 1, 2 });
t.fill(0.0);
const n: usize = t.size(); // total elements
const flat: usize = try t.flatIndex(&.{ 1, 2 }); // offset into t.dataat and setAt check the index count and each bound, returning
ShapeMismatch or OutOfRange. They are the readable path, not the fast
one: for a hot loop, take t.data and the strides and walk them yourself.
These return a new Tensor sharing the same data. They permute or
insert entries in the shape and stride arrays and copy nothing.
const tr = try t.transpose(0, 1); // exchange two axes
const p = try t.permute(&.{ 2, 0, 1 }); // any permutation
const m = try t.moveAxis(0, 2); // move one axis, others keep order
const up = try t.unsqueeze(1); // insert a size-1 axis
const down = try up.squeeze(1); // remove one
const flat = try t.flatten(); // rank-1 view; contiguous only
const r = try t.reshape(&.{ 4, 6 }); // contiguous only, same element countWriting through one is visible through the others: they are the same memory.
permute requires a genuine permutation of 0..rank; a repeated axis is
OutOfRange. squeeze refuses an axis whose extent is not 1, because
removing it would change how many elements the view addresses.
const batch = try data.slice(0, 128, 32); // 32 rows starting at row 128slice keeps every stride and moves base forward by
start × strides[axis], then shortens that axis. Nothing is copied, and slices
compose: a slice of a slice adds the offsets.
const mid = try t.slice(0, 1, 2); // rows 1 and 2
try mid.setAt(&.{ 0, 0 }, -1.0); // visible as t[1][0]start + len must fit; a window that would run past the end is
OutOfRange rather than a shorter view than requested.
var shape: [zn.max_rank]usize = @splat(0);
const rank = try zn.broadcastShape(&.{ 3, 1, 5 }, &.{ 4, 5 }, &shape);
// rank == 3, shape[0..3] == { 3, 4, 5 }
const wide = try row.broadcastTo(&.{ 4, 3 }); // row is (1, 3)
// wide.strides[0] == 0 — all four rows read the same three valuesmax_rank/// The most axes a tensor may have. Shape and strides are fixed-size arrays of this length, so a
/// `Tensor` is a small value with no allocation of its own and copying one is free.
///
/// Six because it covers what this library will actually meet - batch, channel, depth, height,
/// width, and one spare - and because the arrays are indexed in inner loops, where a bound the
/// compiler knows is worth more than generality nobody uses.
pub const max_rank: usize = 6;
broadcastShape/// The broadcast shape of two operands, or `ShapeMismatch` when they do not broadcast.
///
/// The NumPy rule: align the shapes at their TRAILING axis, and for each pair the extents must be
/// equal or one of them must be 1, in which case the other wins. A missing leading axis behaves
/// as 1, so `(3, 1, 5)` and `(4, 5)` broadcast to `(3, 4, 5)`.
///
/// Returns the shape rather than a bool, because every caller that asks whether two shapes
/// broadcast immediately needs the result, and computing it twice invites the two answers to
/// disagree.
pub fn broadcastShape(
a: []const usize,
b: []const usize,
out_shape: *[max_rank]usize,
) Error!u8 {
const rank: usize = @max(a.len, b.len);
if (rank > max_rank) {
return Error.UnsupportedShape;
}
var axis: usize = 0;
while (axis < rank) : (axis += 1) {
// Count from the trailing axis; an axis the shorter operand does not have behaves as 1.
const from_end: usize = rank - axis;
const ea: usize = if (a.len >= from_end) a[a.len - from_end] else 1;
const eb: usize = if (b.len >= from_end) b[b.len - from_end] else 1;
if (ea != eb and ea != 1 and eb != 1) {
return Error.ShapeMismatch;
}
out_shape[axis] = @max(ea, eb);
}
return @intCast(rank);
}
The NumPy rule: align the shapes at their trailing axis; for each pair the extents must be
equal or one must be 1, in which case the other wins. A missing leading axis behaves as 1. Shapes
that do not broadcast return ShapeMismatch — (3) against
(4) is an error, not a guess.
Stretching sets the stride of a size-1 axis to zero, so every position along it addresses the same element. Nothing is allocated and no data moves.
map, zip and everything built
on them return UnsupportedShape when out.isAliased().
Two predicates decide what may be done with a view.
| Predicate | False when | Consequences |
|---|---|---|
isContiguous() | The strides are not dense row-major — after transpose, permute, slice or a stretch | reshape and flatten refuse; anything walking data linearly must not; a GPU kernel cannot bind it |
isAliased() | Never — it is true when a stretched axis has extent above 1 | Every writing operation refuses such a destination |
pub fn isAliased(self: Self) bool {
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
if (self.shape[axis] > 1 and self.strides[axis] == 0) return true;
}
return false;
}try zn.materialise(f32, dense, view); // any view -> dense
try zn.concat(f32, out, a, b, 0); // join along an axis
try zn.tile(f32, out, a, 0); // repeat along an axis
try zn.take(f32, out, a, &.{ 3, 0, 3 }, 0); // gather named positionsmaterialise/// Copy `a` into `out`, element for element. `out` must have `a`'s shape.
///
/// THIS IS `contiguous` IN znum's SENSE, AND IT IS THE DENSIFY STEP
///
/// The view operations above copy nothing, which is what makes them free - and also what makes
/// them unbindable: a kernel takes a dense buffer and indexes from its start. `materialise` is
/// how a permuted, transposed or stretched view becomes something a GPU can read, and every
/// dispatch that meets such a view must call it first.
///
/// It reads through `at`, so ANY view works as the source: strided, permuted, stretched. The
/// destination must not be aliased, for `zip`'s reason.
pub fn materialise(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
return map(T, out, a, struct {
fn apply(x: T) T {
return x;
}
}.apply);
}
concat/// Join `a` and `b` along `axis` into `out`.
///
/// Every axis except `axis` must match, and `out`'s `axis` extent must be the SUM. Checked
/// before anything is written, so a partially-filled output is not a state this can produce.
pub fn concat(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
axis: usize,
) Error!void {
if (a.rank != b.rank or out.rank != a.rank) {
return Error.ShapeMismatch;
}
if (axis >= a.rank) {
return Error.OutOfRange;
}
var i: usize = 0;
while (i < a.rank) : (i += 1) {
const want: usize = if (i == axis) a.shape[i] + b.shape[i] else a.shape[i];
if (out.shape[i] != want) {
return Error.ShapeMismatch;
}
if (i != axis and a.shape[i] != b.shape[i]) {
return Error.ShapeMismatch;
}
}
if (out.isAliased()) {
return Error.UnsupportedShape;
}
// Walk the OUTPUT and pick the source per position: `axis` below the split comes from `a`,
// at or above it from `b` with the offset removed. Walking the output means every element is
// written exactly once by construction.
var walker: [max_rank]usize = @splat(0);
var src: [max_rank]usize = @splat(0);
const shape: []const usize = out.shape[0..out.rank];
var remaining: usize = out.size();
while (remaining > 0) : (remaining -= 1) {
const at: []const usize = walker[0..out.rank];
@memcpy(src[0..out.rank], at);
const from_a: bool = at[axis] < a.shape[axis];
if (!from_a) {
src[axis] = at[axis] - a.shape[axis];
}
const source: Tensor(T) = if (from_a) a else b;
out.data[out.offsetOf(at)] =
source.data[source.offsetOf(src[0..out.rank])];
advance(shape, walker[0..out.rank]);
}
}
tile/// Repeat `a` along `axis` until `out` is filled. `out`'s `axis` extent must be a whole multiple
/// of `a`'s.
///
/// A whole multiple, not "at least": a partial final tile is a different operation, and
/// silently producing one would make `tile` and `pad` indistinguishable at the call site.
pub fn tile(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void {
if (out.rank != a.rank) {
return Error.ShapeMismatch;
}
if (axis >= a.rank or a.shape[axis] == 0) {
return Error.OutOfRange;
}
var i: usize = 0;
while (i < a.rank) : (i += 1) {
if (i == axis) {
if (out.shape[i] % a.shape[i] != 0) {
return Error.ShapeMismatch;
}
} else if (out.shape[i] != a.shape[i]) {
return Error.ShapeMismatch;
}
}
if (out.isAliased()) {
return Error.UnsupportedShape;
}
var walker: [max_rank]usize = @splat(0);
var src: [max_rank]usize = @splat(0);
const shape: []const usize = out.shape[0..out.rank];
var remaining: usize = out.size();
while (remaining > 0) : (remaining -= 1) {
const at: []const usize = walker[0..out.rank];
@memcpy(src[0..out.rank], at);
src[axis] = at[axis] % a.shape[axis];
out.data[out.offsetOf(at)] =
a.data[a.offsetOf(src[0..out.rank])];
advance(shape, walker[0..out.rank]);
}
}
take/// Gather rows of `a` named by `indices` along `axis` into `out`.
///
/// `indices` are `usize` in host memory rather than a `Tensor`, because on the CPU that is what
/// a caller has - a slice of positions. The GPU form will take an index BUFFER, and that
/// difference is real: this signature will not be the kernel's.
///
/// An out-of-range index is `OutOfRange`, checked BEFORE any element is written. Silently
/// clamping - which is what a GPU kernel is tempted to do, having nowhere to report - would fill
/// the output with a plausible wrong row.
pub fn take(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
indices: []const usize,
axis: usize,
) Error!void {
if (out.rank != a.rank) {
return Error.ShapeMismatch;
}
if (axis >= a.rank) {
return Error.OutOfRange;
}
if (out.shape[axis] != indices.len) {
return Error.ShapeMismatch;
}
var i: usize = 0;
while (i < a.rank) : (i += 1) {
if (i != axis and out.shape[i] != a.shape[i]) {
return Error.ShapeMismatch;
}
}
for (indices) |idx| {
if (idx >= a.shape[axis]) {
return Error.OutOfRange;
}
}
if (out.isAliased()) {
return Error.UnsupportedShape;
}
var walker: [max_rank]usize = @splat(0);
var src: [max_rank]usize = @splat(0);
const shape: []const usize = out.shape[0..out.rank];
var remaining: usize = out.size();
while (remaining > 0) : (remaining -= 1) {
const at: []const usize = walker[0..out.rank];
@memcpy(src[0..out.rank], at);
src[axis] = indices[at[axis]];
out.data[out.offsetOf(at)] =
a.data[a.offsetOf(src[0..out.rank])];
advance(shape, walker[0..out.rank]);
}
}
materialise is the step that turns a strided or permuted view into something dense
— needed before a view can be handed to a GPU kernel, which indexes from the start of a
buffer. It reads through at, so any view works as the source.
concat requires every axis except the join axis to match and the output's join
extent to be the sum. tile requires a whole multiple.
take checks every index before writing anything; repeats are allowed and the
output order follows the index list.
const Activations = zn.Named(f64, &.{ "batch", "width" });
const a = try Activations.of(raw);
try zn.sumAxis(f64, out, a.tensor, comptime Activations.axis("width"));
const flipped = try a.transposed("batch", "width"); // names follow the data
try zn.namedMatmul(f64, &.{ "batch", "width" }, &.{ "width", "classes" }, logits, a, w);Named/// A tensor whose axes carry names, checked at compile time.
///
/// THE MAGIC NUMBER IS THE BUG THIS REMOVES
///
/// `sumAxis(out, activations, 1)` is a sentence with a hole in it. Which axis is 1? On a
/// `(batch, width)` tensor it is the width; on a `(width, batch)` one it is the batch, and the
/// two produce different numbers of the same type with no complaint from anything. Every axis
/// argument in this file is that hole, and the only defence has been care.
///
/// `activations.axis("width")` closes it. The names live in the TYPE, so the lookup happens at
/// compile time and costs nothing at runtime - and a name that is not there is a compile error
/// naming it, rather than an index that happens to be in range.
///
/// WHAT THIS IS NOT: it is not a new tensor. `Named` wraps a `Tensor(T)` and hands it back
/// through `.tensor`, so every function in this file keeps working unchanged and a caller adopts
/// names where they help rather than everywhere at once. znum's `dimnames` is 3 302 lines
/// building a parallel API; this is the part that carries the weight.
pub fn Named(comptime T: type, comptime names: []const []const u8) type {
comptime {
if (names.len > max_rank) {
@compileError("Named: too many axes");
}
for (names, 0..) |name, i| {
if (name.len == 0) {
@compileError("Named: an axis name may not be empty");
}
for (names[0..i]) |earlier| {
if (std.mem.eql(u8, earlier, name)) {
@compileError("Named: axis name `" ++ name ++ "` is repeated");
}
}
}
}
return struct {
const Self = @This();
/// The tensor itself. Every zimrnum function takes this directly.
tensor: Tensor(T),
/// The axis names, in order.
pub const axis_names: []const []const u8 = names;
/// Attach names to a tensor. The rank must match, which is the one thing that cannot be
/// checked at compile time because rank is a runtime field.
pub fn of(tensor: Tensor(T)) Error!Self {
if (tensor.rank != names.len) {
return Error.ShapeMismatch;
}
return .{ .tensor = tensor };
}
/// The index of the axis called `name`. A COMPILE ERROR if there is no such axis, which
/// is the whole point: a typo cannot become an in-range index.
pub fn axis(comptime name: []const u8) usize {
comptime {
for (names, 0..) |candidate, i| {
if (std.mem.eql(u8, candidate, name)) {
return i;
}
}
// The message lists what IS there, because "no axis called `wdith`" next to
// "has: batch width" is a typo a reader fixes in a second.
var known: []const u8 = "";
for (names) |candidate| {
known = known ++ candidate ++ " ";
}
@compileError("Named: no axis called `" ++ name ++ "`; has: " ++ known);
}
}
/// The extent of the axis called `name`.
pub fn extent(self: Self, comptime name: []const u8) usize {
return self.tensor.shape[comptime axis(name)];
}
/// Sum along a NAMED axis into `out`, whose names are this one's with that axis removed.
///
/// The reason names exist. `sumAxis(out, a, 1)` and `a.sumAxis(out, "width")` compile to
/// the same loop, and only one of them says which axis it meant.
pub fn sumAxis(
self: Self,
comptime name: []const u8,
out: Named(T, dropName(names, name)),
) Error!void {
return zimrnum.sumAxis(T, out.tensor, self.tensor, comptime axis(name));
}
/// Mean along a named axis.
pub fn meanAxis(
self: Self,
comptime name: []const u8,
out: Named(T, dropName(names, name)),
) Error!void {
return zimrnum.meanAxis(T, out.tensor, self.tensor, comptime axis(name));
}
/// Largest value along a named axis.
pub fn maxAxis(
self: Self,
comptime name: []const u8,
out: Named(T, dropName(names, name)),
) Error!void {
return zimrnum.maxAxis(T, out.tensor, self.tensor, comptime axis(name));
}
/// A window of `len` positions along a named axis, keeping every name.
///
/// Batching is exactly this: `activations.window("batch", 0, 32)` is a `Named` of the
/// same type, so nothing downstream changes and nothing has to be told which axis was
/// cut.
pub fn window(
self: Self,
comptime name: []const u8,
start: usize,
len: usize,
) Error!Self {
return Self.of(try self.tensor.slice(comptime axis(name), start, len));
}
/// This tensor with two axes exchanged, and the names exchanged to match.
///
/// The names follow the data, which is the property that makes the whole thing worth
/// having: a transposed `(batch, width)` is a `(width, batch)`, and every later
/// `axis("width")` on it returns the new index without the caller adjusting anything.
pub fn transposed(
self: Self,
comptime first: []const u8,
comptime second: []const u8,
) Error!Named(T, swapNames(names, first, second)) {
const i: usize = comptime axis(first);
const j: usize = comptime axis(second);
return Named(T, swapNames(names, first, second))
.of(try self.tensor.transpose(i, j));
}
};
}
sumAxis/// Sum `a` along `axis` into `out`, which must have `a`'s shape with that axis removed.
///
/// The output's rank is one less, not one with a size-1 axis. Keeping the axis would be the
/// `keepdims` option, and options that change the RESULT SHAPE are better as separate calls than
/// as a flag a reader has to look up - a caller who wants the axis back has `reshape`.
pub fn sumAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void {
comptime requireNumeric(T);
if (axis >= a.rank) {
return Error.OutOfRange;
}
if (out.rank + 1 != a.rank or out.isAliased()) {
return Error.ShapeMismatch;
}
// The surviving axes must match, in order.
var src: usize = 0;
var dst: usize = 0;
while (src < a.rank) : (src += 1) {
if (src == axis) {
continue;
}
if (a.shape[src] != out.shape[dst]) {
return Error.ShapeMismatch;
}
dst += 1;
}
out.fill(0);
var walker: [max_rank]usize = @splat(0);
var reduced: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
var remaining: usize = a.size();
while (remaining > 0) : (remaining -= 1) {
// Project the source position onto the output by dropping the reduced axis.
var k: usize = 0;
var i: usize = 0;
while (i < a.rank) : (i += 1) {
if (i == axis) {
continue;
}
reduced[k] = walker[i];
k += 1;
}
const at: []const usize = reduced[0..out.rank];
const slot: usize = out.offsetOf(at);
out.data[slot] += a.data[a.offsetOf(walker[0..a.rank])];
advance(shape, walker[0..a.rank]);
}
}
namedMatmul/// `out = a @ b`, with the contracted axes required to have THE SAME NAME.
///
/// THE CHECK THAT SHAPES CANNOT MAKE
///
/// `matmul` verifies that `a`'s columns equal `b`'s rows - a check on NUMBERS. Two tensors can
/// agree on that and still be the wrong pair: a `(batch, width)` times a `(width, classes)` is a
/// layer, and a `(batch, width)` times a `(batch, classes)` is nonsense that happens to typecheck
/// whenever batch and width are equal. Nothing in this file could tell them apart before.
///
/// Requiring the contracted axes to share a name is a check on MEANING, and it happens at compile
/// time. The result is named for the two surviving axes.
pub fn namedMatmul(
comptime T: type,
comptime a_names: []const []const u8,
comptime b_names: []const []const u8,
out: Named(T, &.{ a_names[0], b_names[1] }),
a: Named(T, a_names),
b: Named(T, b_names),
) Error!void {
comptime {
if (a_names.len != 2 or b_names.len != 2) {
@compileError("namedMatmul: both operands must be rank 2");
}
if (!std.mem.eql(u8, a_names[1], b_names[0])) {
@compileError("namedMatmul: contracting `" ++ a_names[1] ++ "` against `" ++
b_names[0] ++ "` - the axes being summed over must be the same axis");
}
}
return matmul(T, out.tensor, a.tensor, b.tensor);
}
sumAxis(out, activations, 1) is a sentence with a hole in it. Which axis is 1? On
a (batch, width) tensor it is the width; on a (width, batch) one it is
the batch, and the two produce different numbers of the same type with no complaint from anything.
Every axis argument in this library is that hole, and the only defence has been
care.
The names live in the type, so axis("width") resolves at compile time and
costs nothing at runtime — and a name that is not there is a compile error naming it:
| Written | Result |
|---|---|
A.axis("wdith") | Named: no axis called `wdith`; has: batch width |
Named(f64, &.{"batch","batch"}) | Named: axis name `batch` is repeated |
contracting width against batch | namedMatmul: the axes being summed over must be the same axis |
matmul checks that a's columns equal b's rows — a
check on numbers. Two tensors can agree on that and still be the wrong pair: a
(batch, width) times a (width, classes) is a layer, and a
(batch, width) times a (batch, classes) is nonsense that typechecks
whenever batch and width happen to be equal. namedMatmul checks
meaning instead, at compile time.
The names follow the data. A transposed (batch, width) is a
(width, batch), and every later axis("width") on it returns the new index
without the caller adjusting anything. That is the property that makes names worth having rather
than merely pleasant.
The reductions take a name instead of an index, and the result's type has that axis removed:
const Logits = zn.Named(f64, &.{ "batch", "classes" });
const PerClass = zn.Named(f64, &.{"classes"});
try logits.meanAxis("batch", per_class); // averages over the batch
try logits.maxAxis("classes", per_sample); // one number per sample
const mini = try inputs.window("batch", 2, 3);Named/// A tensor whose axes carry names, checked at compile time.
///
/// THE MAGIC NUMBER IS THE BUG THIS REMOVES
///
/// `sumAxis(out, activations, 1)` is a sentence with a hole in it. Which axis is 1? On a
/// `(batch, width)` tensor it is the width; on a `(width, batch)` one it is the batch, and the
/// two produce different numbers of the same type with no complaint from anything. Every axis
/// argument in this file is that hole, and the only defence has been care.
///
/// `activations.axis("width")` closes it. The names live in the TYPE, so the lookup happens at
/// compile time and costs nothing at runtime - and a name that is not there is a compile error
/// naming it, rather than an index that happens to be in range.
///
/// WHAT THIS IS NOT: it is not a new tensor. `Named` wraps a `Tensor(T)` and hands it back
/// through `.tensor`, so every function in this file keeps working unchanged and a caller adopts
/// names where they help rather than everywhere at once. znum's `dimnames` is 3 302 lines
/// building a parallel API; this is the part that carries the weight.
pub fn Named(comptime T: type, comptime names: []const []const u8) type {
comptime {
if (names.len > max_rank) {
@compileError("Named: too many axes");
}
for (names, 0..) |name, i| {
if (name.len == 0) {
@compileError("Named: an axis name may not be empty");
}
for (names[0..i]) |earlier| {
if (std.mem.eql(u8, earlier, name)) {
@compileError("Named: axis name `" ++ name ++ "` is repeated");
}
}
}
}
return struct {
const Self = @This();
/// The tensor itself. Every zimrnum function takes this directly.
tensor: Tensor(T),
/// The axis names, in order.
pub const axis_names: []const []const u8 = names;
/// Attach names to a tensor. The rank must match, which is the one thing that cannot be
/// checked at compile time because rank is a runtime field.
pub fn of(tensor: Tensor(T)) Error!Self {
if (tensor.rank != names.len) {
return Error.ShapeMismatch;
}
return .{ .tensor = tensor };
}
/// The index of the axis called `name`. A COMPILE ERROR if there is no such axis, which
/// is the whole point: a typo cannot become an in-range index.
pub fn axis(comptime name: []const u8) usize {
comptime {
for (names, 0..) |candidate, i| {
if (std.mem.eql(u8, candidate, name)) {
return i;
}
}
// The message lists what IS there, because "no axis called `wdith`" next to
// "has: batch width" is a typo a reader fixes in a second.
var known: []const u8 = "";
for (names) |candidate| {
known = known ++ candidate ++ " ";
}
@compileError("Named: no axis called `" ++ name ++ "`; has: " ++ known);
}
}
/// The extent of the axis called `name`.
pub fn extent(self: Self, comptime name: []const u8) usize {
return self.tensor.shape[comptime axis(name)];
}
/// Sum along a NAMED axis into `out`, whose names are this one's with that axis removed.
///
/// The reason names exist. `sumAxis(out, a, 1)` and `a.sumAxis(out, "width")` compile to
/// the same loop, and only one of them says which axis it meant.
pub fn sumAxis(
self: Self,
comptime name: []const u8,
out: Named(T, dropName(names, name)),
) Error!void {
return zimrnum.sumAxis(T, out.tensor, self.tensor, comptime axis(name));
}
/// Mean along a named axis.
pub fn meanAxis(
self: Self,
comptime name: []const u8,
out: Named(T, dropName(names, name)),
) Error!void {
return zimrnum.meanAxis(T, out.tensor, self.tensor, comptime axis(name));
}
/// Largest value along a named axis.
pub fn maxAxis(
self: Self,
comptime name: []const u8,
out: Named(T, dropName(names, name)),
) Error!void {
return zimrnum.maxAxis(T, out.tensor, self.tensor, comptime axis(name));
}
/// A window of `len` positions along a named axis, keeping every name.
///
/// Batching is exactly this: `activations.window("batch", 0, 32)` is a `Named` of the
/// same type, so nothing downstream changes and nothing has to be told which axis was
/// cut.
pub fn window(
self: Self,
comptime name: []const u8,
start: usize,
len: usize,
) Error!Self {
return Self.of(try self.tensor.slice(comptime axis(name), start, len));
}
/// This tensor with two axes exchanged, and the names exchanged to match.
///
/// The names follow the data, which is the property that makes the whole thing worth
/// having: a transposed `(batch, width)` is a `(width, batch)`, and every later
/// `axis("width")` on it returns the new index without the caller adjusting anything.
pub fn transposed(
self: Self,
comptime first: []const u8,
comptime second: []const u8,
) Error!Named(T, swapNames(names, first, second)) {
const i: usize = comptime axis(first);
const j: usize = comptime axis(second);
return Named(T, swapNames(names, first, second))
.of(try self.tensor.transpose(i, j));
}
};
}
Naming the result is what makes this more than sugar. meanAxis("batch", x) where
x is still named (batch, classes) does not compile — the axis is
gone from the type, so a result of the wrong shape cannot be passed on. And
window returns the same type, which is what a mini-batch is: fewer rows, same
meaning, nothing downstream needing to be told.
A two-layer network reads as a sentence, and every transition in it is checked at compile time:
const Inputs = zn.Named(f64, &.{ "batch", "features" });
const W1 = zn.Named(f64, &.{ "features", "hidden" });
const Acts = zn.Named(f64, &.{ "batch", "hidden" });
const W2 = zn.Named(f64, &.{ "hidden", "classes" });
const Logits = zn.Named(f64, &.{ "batch", "classes" });
try zn.namedMatmul(f64, Inputs.axis_names, W1.axis_names, acts, inputs, w1);
try zn.relu(f64, acts.tensor, acts.tensor);
try zn.namedMatmul(f64, Acts.axis_names, W2.axis_names, logits, acts, w2);Named/// A tensor whose axes carry names, checked at compile time.
///
/// THE MAGIC NUMBER IS THE BUG THIS REMOVES
///
/// `sumAxis(out, activations, 1)` is a sentence with a hole in it. Which axis is 1? On a
/// `(batch, width)` tensor it is the width; on a `(width, batch)` one it is the batch, and the
/// two produce different numbers of the same type with no complaint from anything. Every axis
/// argument in this file is that hole, and the only defence has been care.
///
/// `activations.axis("width")` closes it. The names live in the TYPE, so the lookup happens at
/// compile time and costs nothing at runtime - and a name that is not there is a compile error
/// naming it, rather than an index that happens to be in range.
///
/// WHAT THIS IS NOT: it is not a new tensor. `Named` wraps a `Tensor(T)` and hands it back
/// through `.tensor`, so every function in this file keeps working unchanged and a caller adopts
/// names where they help rather than everywhere at once. znum's `dimnames` is 3 302 lines
/// building a parallel API; this is the part that carries the weight.
pub fn Named(comptime T: type, comptime names: []const []const u8) type {
comptime {
if (names.len > max_rank) {
@compileError("Named: too many axes");
}
for (names, 0..) |name, i| {
if (name.len == 0) {
@compileError("Named: an axis name may not be empty");
}
for (names[0..i]) |earlier| {
if (std.mem.eql(u8, earlier, name)) {
@compileError("Named: axis name `" ++ name ++ "` is repeated");
}
}
}
}
return struct {
const Self = @This();
/// The tensor itself. Every zimrnum function takes this directly.
tensor: Tensor(T),
/// The axis names, in order.
pub const axis_names: []const []const u8 = names;
/// Attach names to a tensor. The rank must match, which is the one thing that cannot be
/// checked at compile time because rank is a runtime field.
pub fn of(tensor: Tensor(T)) Error!Self {
if (tensor.rank != names.len) {
return Error.ShapeMismatch;
}
return .{ .tensor = tensor };
}
/// The index of the axis called `name`. A COMPILE ERROR if there is no such axis, which
/// is the whole point: a typo cannot become an in-range index.
pub fn axis(comptime name: []const u8) usize {
comptime {
for (names, 0..) |candidate, i| {
if (std.mem.eql(u8, candidate, name)) {
return i;
}
}
// The message lists what IS there, because "no axis called `wdith`" next to
// "has: batch width" is a typo a reader fixes in a second.
var known: []const u8 = "";
for (names) |candidate| {
known = known ++ candidate ++ " ";
}
@compileError("Named: no axis called `" ++ name ++ "`; has: " ++ known);
}
}
/// The extent of the axis called `name`.
pub fn extent(self: Self, comptime name: []const u8) usize {
return self.tensor.shape[comptime axis(name)];
}
/// Sum along a NAMED axis into `out`, whose names are this one's with that axis removed.
///
/// The reason names exist. `sumAxis(out, a, 1)` and `a.sumAxis(out, "width")` compile to
/// the same loop, and only one of them says which axis it meant.
pub fn sumAxis(
self: Self,
comptime name: []const u8,
out: Named(T, dropName(names, name)),
) Error!void {
return zimrnum.sumAxis(T, out.tensor, self.tensor, comptime axis(name));
}
/// Mean along a named axis.
pub fn meanAxis(
self: Self,
comptime name: []const u8,
out: Named(T, dropName(names, name)),
) Error!void {
return zimrnum.meanAxis(T, out.tensor, self.tensor, comptime axis(name));
}
/// Largest value along a named axis.
pub fn maxAxis(
self: Self,
comptime name: []const u8,
out: Named(T, dropName(names, name)),
) Error!void {
return zimrnum.maxAxis(T, out.tensor, self.tensor, comptime axis(name));
}
/// A window of `len` positions along a named axis, keeping every name.
///
/// Batching is exactly this: `activations.window("batch", 0, 32)` is a `Named` of the
/// same type, so nothing downstream changes and nothing has to be told which axis was
/// cut.
pub fn window(
self: Self,
comptime name: []const u8,
start: usize,
len: usize,
) Error!Self {
return Self.of(try self.tensor.slice(comptime axis(name), start, len));
}
/// This tensor with two axes exchanged, and the names exchanged to match.
///
/// The names follow the data, which is the property that makes the whole thing worth
/// having: a transposed `(batch, width)` is a `(width, batch)`, and every later
/// `axis("width")` on it returns the new index without the caller adjusting anything.
pub fn transposed(
self: Self,
comptime first: []const u8,
comptime second: []const u8,
) Error!Named(T, swapNames(names, first, second)) {
const i: usize = comptime axis(first);
const j: usize = comptime axis(second);
return Named(T, swapNames(names, first, second))
.of(try self.tensor.transpose(i, j));
}
};
}
namedMatmul/// `out = a @ b`, with the contracted axes required to have THE SAME NAME.
///
/// THE CHECK THAT SHAPES CANNOT MAKE
///
/// `matmul` verifies that `a`'s columns equal `b`'s rows - a check on NUMBERS. Two tensors can
/// agree on that and still be the wrong pair: a `(batch, width)` times a `(width, classes)` is a
/// layer, and a `(batch, width)` times a `(batch, classes)` is nonsense that happens to typecheck
/// whenever batch and width are equal. Nothing in this file could tell them apart before.
///
/// Requiring the contracted axes to share a name is a check on MEANING, and it happens at compile
/// time. The result is named for the two surviving axes.
pub fn namedMatmul(
comptime T: type,
comptime a_names: []const []const u8,
comptime b_names: []const []const u8,
out: Named(T, &.{ a_names[0], b_names[1] }),
a: Named(T, a_names),
b: Named(T, b_names),
) Error!void {
comptime {
if (a_names.len != 2 or b_names.len != 2) {
@compileError("namedMatmul: both operands must be rank 2");
}
if (!std.mem.eql(u8, a_names[1], b_names[0])) {
@compileError("namedMatmul: contracting `" ++ a_names[1] ++ "` against `" ++
b_names[0] ++ "` - the axes being summed over must be the same axis");
}
}
return matmul(T, out.tensor, a.tensor, b.tensor);
}
relu/// `out = max(a, 0)`, elementwise.
pub fn relu(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return if (x > 0) x else 0;
}
}.apply);
}
Where names do not reach: inside the autograd tape. graph.matmul
returns a Var, which is an index into the graph and carries no names. Names help
where you hold tensors and pass axis indices — reductions, products, windows — and the
tape is deliberately not that. Saying so is better than implying a guarantee that stops at the
first graph. call.
Named wraps a Tensor and hands it back through .tensor,
so every function here keeps working unchanged and names are adopted where they help rather than
everywhere at once. znum's dimnames is 3 302 lines building a parallel API; this
is the part that carries the weight.
try zn.meshGrid(f64, grid_x, grid_y, xs, ys); // both grids, both NAMED
try zn.repeatEach(f64, out, a, 0, 2); // [1,2,3] -> [1,1,2,2,3,3]
try zn.tile(f64, out, a, 0); // [1,2,3] -> [1,2,3,1,2,3]
const moved = try zn.moveAxis(f64, cube, 0, 2); // a view, costs nothingmeshGrid/// The two coordinate grids of a rectangular mesh.
///
/// `out_x[i][j] = x[j]` and `out_y[i][j] = y[i]`, so a pair of 1-D axes becomes the full grid of
/// positions a surface or a plot needs. `x` supplies the columns and `y` the rows, which is the
/// order every plotting library uses and the opposite of what the argument order suggests - hence
/// the outputs being named rather than returned as a two-element array.
///
/// **znum returns `[2]Tensor(T)` and the caller has to remember which is which.** Here the caller
/// allocates both and names them at the call site, so `out_x` cannot be mistaken for `out_y`.
pub fn meshGrid(
comptime T: type,
out_x: Tensor(T),
out_y: Tensor(T),
x: Tensor(T),
y: Tensor(T),
) Error!void {
comptime requireNumeric(T);
if (x.rank != 1 or y.rank != 1) {
return Error.UnsupportedShape;
}
const rows: usize = y.shape[0];
const cols: usize = x.shape[0];
if (out_x.rank != 2 or out_y.rank != 2) {
return Error.UnsupportedShape;
}
if (out_x.shape[0] != rows or out_x.shape[1] != cols) {
return Error.ShapeMismatch;
}
if (out_y.shape[0] != rows or out_y.shape[1] != cols) {
return Error.ShapeMismatch;
}
for (0..rows) |row| {
for (0..cols) |col| {
try out_x.setAt(&.{ row, col }, try x.at(&.{col}));
try out_y.setAt(&.{ row, col }, try y.at(&.{row}));
}
}
}
repeatEach/// Each element repeated `count` times along `axis`.
///
/// REPEATING ELEMENTS IS NOT REPEATING THE TENSOR, AND THE NAMES SHOULD SAY WHICH
///
/// numpy calls this `repeat` and calls the other one `tile`, and the pair is a well-known source
/// of confusion because neither name says what varies:
///
/// repeatEach([1, 2, 3], 2) -> [1, 1, 2, 2, 3, 3] each element, in place
/// tile ([1, 2, 3], 2) -> [1, 2, 3, 1, 2, 3] the whole thing, end to end
///
/// `tile` was already here under numpy's name; this one is `repeatEach` rather than `repeat`
/// because the "each" is the entire difference.
pub fn repeatEach(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
axis: usize,
count: usize,
) Error!void {
comptime requireNumeric(T);
if (out.rank != a.rank or axis >= a.rank) {
return Error.UnsupportedShape;
}
if (count == 0) {
return Error.DomainError;
}
for (0..a.rank) |d| {
const want: usize = if (d == axis) a.shape[d] * count else a.shape[d];
if (out.shape[d] != want) {
return Error.ShapeMismatch;
}
}
var source: [max_rank]usize = @splat(0);
var walk: Walk = .over(out.shape[0..out.rank]);
while (walk.next()) |at| {
for (0..out.rank) |d| {
// Every position maps back by dividing the repeated axis; the rest pass through.
source[d] = if (d == axis) at[d] / count else at[d];
}
try out.setAt(at, try a.at(source[0..a.rank]));
}
}
tile/// Repeat `a` along `axis` until `out` is filled. `out`'s `axis` extent must be a whole multiple
/// of `a`'s.
///
/// A whole multiple, not "at least": a partial final tile is a different operation, and
/// silently producing one would make `tile` and `pad` indistinguishable at the call site.
pub fn tile(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void {
if (out.rank != a.rank) {
return Error.ShapeMismatch;
}
if (axis >= a.rank or a.shape[axis] == 0) {
return Error.OutOfRange;
}
var i: usize = 0;
while (i < a.rank) : (i += 1) {
if (i == axis) {
if (out.shape[i] % a.shape[i] != 0) {
return Error.ShapeMismatch;
}
} else if (out.shape[i] != a.shape[i]) {
return Error.ShapeMismatch;
}
}
if (out.isAliased()) {
return Error.UnsupportedShape;
}
var walker: [max_rank]usize = @splat(0);
var src: [max_rank]usize = @splat(0);
const shape: []const usize = out.shape[0..out.rank];
var remaining: usize = out.size();
while (remaining > 0) : (remaining -= 1) {
const at: []const usize = walker[0..out.rank];
@memcpy(src[0..out.rank], at);
src[axis] = at[axis] % a.shape[axis];
out.data[out.offsetOf(at)] =
a.data[a.offsetOf(src[0..out.rank])];
advance(shape, walker[0..out.rank]);
}
}
moveAxis/// A view of `a` with `axis` moved to sit at `destination`, the others closing the gap.
///
/// A VIEW, NOT A COPY - AND THAT IS WHY IT IS NOT `transposed`
///
/// `transposed(i, j)` swaps two axes and leaves the rest alone. Moving an axis is different: the
/// axes between source and destination all SHIFT to make room, which is what you want when a
/// batch dimension has to travel from the front to the back. numpy calls it `moveaxis`; this is
/// the same operation with the house's camelCase.
///
/// Like `transposed` it only permutes strides, so it costs nothing and aliases `a`'s storage.
pub fn moveAxis(
comptime T: type,
a: Tensor(T),
axis: usize,
destination: usize,
) Error!Tensor(T) {
if (axis >= a.rank or destination >= a.rank) {
return Error.OutOfRange;
}
var order: [max_rank]usize = @splat(0);
var next: usize = 0;
// Lay down every axis except the one moving, then insert it at `destination`.
for (0..a.rank) |slot| {
if (slot == destination) {
order[slot] = axis;
continue;
}
if (next == axis) {
next += 1;
}
order[slot] = next;
next += 1;
}
return a.permute(order[0..a.rank]);
}
numpy calls the second one repeat and the third tile, and the
pair is a well-known source of confusion because neither name says what varies. Here the
first is repeatEach: the "each" is the difference, and putting it in the name
means the two lines above read differently at a glance instead of needing the docs.
meshGrid takes both outputs as named parameters rather than returning
[2]Tensor. znum returns the pair and the caller has to remember which is which;
grid_x and grid_y cannot be mistaken for each other at the call site.
Note the axis order: x supplies the columns and y the rows,
which is what every plotting library expects and the opposite of what reading the arguments left
to right suggests.
moveAxis is not a swap. Swapping trades two axes and leaves the rest alone; moving
an axis makes the ones in between shift to close the gap — which is what you want when a
batch dimension has to travel from the front to the back. On rank 2 the two are identical, so the
test uses rank 3 where they differ. Like every other view here it only permutes strides, so it
shares storage and costs nothing.
Three of znum's manipulation functions are deliberately absent. hstack,
vstack and dstack are concat with the axis baked in, and
znum's own doc admits the trap ("this does not promote 1-D to 2-D the way numpy's vstack
does") — concat(axis) says the same thing without the footnote. And
reshapeInfer, which lets one dimension be -1, exists because znum
allocates its result and must be told the shape. Here the caller allocates, so the shape
is already known and there is nothing to infer.
resolveInferredShape(12, &.{ 2, -1 }) -> 2 x 6"Two rows, work out the columns" — the one dimension you did not want to compute by hand.
The alternative is x.size() / 2 written at the call site, which is the same arithmetic
with a chance of going stale when the shape changes.
{ -1, -1 } infinitely many answers - picking one invents an intention
{ 2, -2 } -2 is not shorthand for anything; it is a typo for 2 often enough
{ 3, -1 } on 10 elements no integer answer - rounding would drop an elementA shape with no -1 passes through unchanged, so one function serves both cases
— but it still has to account for every element. A shape holding fewer is a silent data loss
and gets the same refusal.
These are the two that get mixed up, because both take a list and an axis and the shapes look alike:
concat two (2,3) along axis 0 -> (4,3) same rank, one axis GREW
stack two (2,3) at axis 0 -> (2,2,3) one rank HIGHER, a NEW axisStacking is "these samples are one batch". Concatenating is "these rows are more rows".
The axis is a parameter rather than always zero because stacking at a later axis interleaves rather than appending — which is what you want when the new dimension belongs inside the existing structure rather than outside it.
The strides trade places and no element moves. So it is cheap, the result may be
non-contiguous, and writing through it writes through to the original — which the
test checks, because a "copy" that turns out to be a view is a bug that surfaces far from its
cause. materialise if you need it packed.
try zn.map(f64, out, src, struct {
fn apply(v: f64) f64 { return zm.sigmoid(v); }
}.apply);
try zn.zip(f64, out, a, b, struct {
fn apply(x: f64, y: f64) f64 { return @max(x, y); }
}.apply);map/// Write `f(x)` into `out` for each element `x` of `a`, broadcasting `a` to `out`'s shape.
pub fn map(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
comptime f: fn (T) T,
) Error!void {
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const va: Tensor(T) = try a.broadcastTo(shape);
const count: usize = out.size();
if (out.isContiguous() and va.isContiguous() and va.data.len >= va.base + count) {
for (out.data[out.base..][0..count], va.data[va.base..][0..count]) |*slot, x| {
slot.* = f(x);
}
return;
}
var walker: [max_rank]usize = @splat(0);
var remaining: usize = count;
while (remaining > 0) : (remaining -= 1) {
const at: []const usize = walker[0..out.rank];
const x: T = va.data[va.offsetOf(at)];
out.data[out.offsetOf(at)] = f(x);
advance(shape, walker[0..out.rank]);
}
}
zip/// Write `f(x, y)` into `out`, broadcasting `a` and `b` to `out`'s shape.
///
/// `out` DEFINES THE SHAPE, AND THE INPUTS MUST REACH IT
///
/// The destination is not inferred. A caller sizes `out` - usually from `broadcastShape` - and
/// both inputs must broadcast to exactly that. Inferring it instead would mean an output whose
/// shape depends on the inputs' shapes in a way the call site does not show, and a mistake there
/// allocates the wrong size rather than reporting anything.
///
/// The contiguous case walks three slices linearly; anything strided or stretched walks the
/// shared index. The fast path additionally checks `data.len`, because a stretched view's `size`
/// exceeds its storage and reasoning about when that cannot happen is worth less than the compare.
pub fn zip(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
comptime f: fn (T, T) T,
) Error!void {
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const va: Tensor(T) = try a.broadcastTo(shape);
const vb: Tensor(T) = try b.broadcastTo(shape);
const count: usize = out.size();
const all_dense: bool = out.isContiguous() and va.isContiguous() and vb.isContiguous() and
va.data.len >= va.base + count and vb.data.len >= vb.base + count;
if (all_dense) {
for (out.data[out.base..][0..count], va.data[va.base..][0..count], vb.data[vb.base..][0..count]) |*slot, x, y| {
slot.* = f(x, y);
}
return;
}
var walker: [max_rank]usize = @splat(0);
var remaining: usize = count;
while (remaining > 0) : (remaining -= 1) {
const at: []const usize = walker[0..out.rank];
const x: T = va.data[va.offsetOf(at)];
const y: T = vb.data[vb.offsetOf(at)];
out.data[out.offsetOf(at)] = f(x, y);
advance(shape, walker[0..out.rank]);
}
}
The function is a comptime parameter, so it inlines and there is no call per element. Every named operation below is a one-line wrapper over one of these.
The output defines the shape. It is not inferred: you size out
— usually from broadcastShape — and both inputs must broadcast to exactly
that, or the call returns ShapeMismatch before writing anything.
Both take a linear pass when every tensor involved is contiguous and an index walk otherwise, so a transposed destination works and gives the same answer.
try zn.add(f64, out, a, b);
try zn.sub(f64, out, a, b);
try zn.mul(f64, out, a, b); // elementwise, not a matrix product
try zn.div(f64, out, a, b); // float only
try zn.scale(f64, out, a, 0.1);
try zn.minimum(f64, out, a, b);
try zn.maximum(f64, out, a, b);
try zn.clamp(f64, out, a, -0.5, 0.5);
try zn.lerp(f64, out, a, b, 0.25);add/// `out = a + b`, broadcasting both to `out`'s shape.
pub fn add(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return x + y;
}
}.apply);
}
sub/// `out = a - b`, broadcasting both to `out`'s shape.
pub fn sub(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return x - y;
}
}.apply);
}
mul/// `out = a * b` ELEMENTWISE, broadcasting both to `out`'s shape. This is not a matrix product.
pub fn mul(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return x * y;
}
}.apply);
}
div/// `out = a / b`, broadcasting both to `out`'s shape.
///
/// Float only. Zig has no single `/` for signed integers - the caller must choose between
/// truncating, flooring and exact division - and picking one here would be this library deciding
/// something the language deliberately asks about. Integer division goes through `zip` with the
/// `@div*` the caller means.
pub fn div(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
comptime requireFloat(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return x / y;
}
}.apply);
}
scale/// `out = a * factor`, elementwise.
///
/// A separate entry point rather than `mul` against a broadcast scalar tensor, because there is
/// no rank-0 tensor here and wrapping one value in a `[1]` allocation to multiply by it is more
/// machinery than the operation. `zip` cannot help either: its function is a comptime parameter
/// and cannot close over a runtime `factor`.
pub fn scale(comptime T: type, out: Tensor(T), a: Tensor(T), factor: T) Error!void {
comptime requireNumeric(T);
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const va: Tensor(T) = try a.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
out.data[out.offsetOf(at)] =
va.data[va.offsetOf(at)] * factor;
}
}
minimum/// `out = min(a, b)`, elementwise. `minimum`, not `min`, because `min` is a reserved math word.
pub fn minimum(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
) Error!void {
comptime requireNumeric(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return @min(x, y);
}
}.apply);
}
maximum/// `out = max(a, b)`, elementwise.
pub fn maximum(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
) Error!void {
comptime requireNumeric(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return @max(x, y);
}
}.apply);
}
clamp/// `out = min(max(a, lo), hi)`, elementwise.
///
/// `lo > hi` is not rejected: the result is then `hi` everywhere, which is what the nested
/// min/max gives and what every other library does. Erroring would be defensible; being SILENTLY
/// different from everyone else would not.
pub fn clamp(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
lo: T,
hi: T,
) Error!void {
comptime requireNumeric(T);
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const va: Tensor(T) = try a.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const x: T = va.data[va.offsetOf(at)];
out.data[out.offsetOf(at)] = @min(@max(x, lo), hi);
}
}
lerp/// `out = a + t * (b - a)`, elementwise.
///
/// NEITHER ONE-LINE FORM IS EXACT AT BOTH ENDS, SO THIS USES TWO
///
/// `a + t*(b-a)` is exact at t = 0 and **not** at t = 1: `x + (y - x)` is not `y` in floating
/// point once the magnitudes differ. `(1-t)*a + t*b` is exact at t = 1 and not at t = 0. I wrote
/// the first form with a comment claiming both, and the endpoint test caught it within a minute.
///
/// So it switches at the midpoint and interpolates from the NEARER endpoint, which is what
/// C++20's `std::lerp` does and for this reason. Exact at both ends, and monotonic in between.
///
/// The branch is on `t` - a scalar the caller passes, identical in every lane - so it costs a
/// GPU nothing. A branch on the DATA would be a different matter.
pub fn lerp(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
t: T,
) Error!void {
comptime requireFloat(T);
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const va: Tensor(T) = try a.broadcastTo(shape);
const vb: Tensor(T) = try b.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const x: T = va.data[va.offsetOf(at)];
const y: T = vb.data[vb.offsetOf(at)];
const v: T = if (t <= 0.5) x + t * (y - x) else y - (1.0 - t) * (y - x);
out.data[out.offsetOf(at)] = v;
}
}
All broadcast their inputs to out's shape, so a row and a column produce a
grid:
const row = try zn.Tensor(f64).alloc(a, &.{ 1, 3 }); // 1, 2, 3
const col = try zn.Tensor(f64).alloc(a, &.{ 4, 1 }); // 10, 20, 30, 40
const out = try zn.Tensor(f64).alloc(a, &.{ 4, 3 });
try zn.add(f64, out, row, col); // out[i][j] == col[i] + row[j]Tensor/// A view of `T`-typed storage as a multidimensional array.
///
/// A TENSOR NEVER OWNS ITS STORAGE, AND THAT IS THE DESIGN
///
/// There is no ownership flag and no `deinit`. `data` is a slice the tensor addresses; whoever
/// allocated it frees it, with the allocator that allocated it. Usually that is an arena, in
/// which case nobody frees anything and `arena.deinit()` ends the matter.
///
/// The alternative - an `owns: bool` and a `deinit` that is a no-op for views - is the shape
/// this library was warned about before it had a tensor: both of its failure modes are SILENT.
/// `deinit` on a view does nothing and reads as if it did something; a missing `deinit` on an
/// owner just grows memory. Having no `deinit` at all means a free appears exactly where an
/// allocation appears, and a view has neither.
///
/// Strides are in ELEMENTS and signed, so a reversed or negatively-strided view is expressible
/// without changing this type. `at` computes a signed offset and casts once.
pub fn Tensor(comptime T: type) type {
return struct {
const Self = @This();
/// The storage this tensor addresses. Not owned - see the type's doc comment.
data: []T,
/// Extent of each axis. Only the first `rank` entries are meaningful.
shape: [max_rank]usize,
/// Step in ELEMENTS between consecutive positions along each axis.
strides: [max_rank]isize,
rank: u8,
/// WHERE THIS VIEW STARTS INSIDE `data`
///
/// Every view until now began at element 0, which is why `slice` could not exist: a
/// window into the middle of a tensor IS an offset. With this, `slice` and the rest of
/// znum's `index` namespace become views rather than copies.
///
/// IT ALSO WEAKENS WHAT `data.len` MEANS. `data` is the whole backing store; `base`
/// and the strides say which of it this view addresses. Everything that walked
/// `data[0..size()]` assuming those coincide had to be corrected - five sites, all of
/// them fast paths that skip the index walk, which is exactly where such an assumption
/// hides.
base: usize = 0,
/// A dense row-major tensor over `shape`, allocated from `gpa`.
///
/// The caller frees `result.data` with the same allocator - or allocates from an arena
/// and frees nothing.
pub fn alloc(gpa: Allocator, shape: []const usize) Error!Self {
if (shape.len > max_rank) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
const storage: []T = gpa.alloc(T, count) catch return Error.OutOfMemory;
return fromSlice(storage, shape);
}
/// A dense row-major view of `data` as `shape`. `data.len` must be exactly the product of
/// the extents - a slice that is merely LARGE ENOUGH is rejected, because the usual reason
/// for the mismatch is a shape that is wrong rather than storage that is generous.
pub fn fromSlice(data: []T, shape: []const usize) Error!Self {
if (shape.len > max_rank) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
if (count != data.len) {
return Error.ShapeMismatch;
}
var self: Self = .{
.data = data,
.shape = @splat(1),
.strides = @splat(0),
.rank = @intCast(shape.len),
.base = 0,
};
// Row-major: the last axis is contiguous, and each earlier stride is the product of
// the extents to its right.
var stride: isize = 1;
var axis: usize = shape.len;
while (axis > 0) {
axis -= 1;
self.shape[axis] = shape[axis];
self.strides[axis] = stride;
stride *= @intCast(shape[axis]);
}
return self;
}
/// The number of elements this tensor addresses.
pub fn size(self: Self) usize {
var count: usize = 1;
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
count *= self.shape[axis];
}
return count;
}
/// Whether the strides are dense row-major, so `data[0..size()]` is the tensor in order.
///
/// An operation that walks memory linearly must ask this rather than assume it: a view
/// produced by `transpose` addresses the same storage in a different order, and a linear
/// walk over it silently reads the wrong elements rather than failing.
pub fn isContiguous(self: Self) bool {
var expected: isize = 1;
var axis: usize = self.rank;
while (axis > 0) {
axis -= 1;
if (self.shape[axis] != 1 and self.strides[axis] != expected) {
return false;
}
expected *= @intCast(self.shape[axis]);
}
return true;
}
/// The offset into `data` of the element at `indices`.
pub fn flatIndex(self: Self, indices: []const usize) Error!usize {
if (indices.len != self.rank) {
return Error.ShapeMismatch;
}
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
if (i >= self.shape[axis]) {
return Error.OutOfRange;
}
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}
/// The element at `indices`.
pub fn at(self: Self, indices: []const usize) Error!T {
return self.data[try self.flatIndex(indices)];
}
/// The flat offset for a coordinate, with no error.
///
/// THE 45 `flatIndex(...) catch unreachable` SITES ARE ALL THIS SHAPE
///
/// const at: []const usize = walker[0..out.rank];
/// out.data[out.offsetOf(at)] = ...
///
/// The slice is `walker[0..rank]`, so `indices.len != rank` **cannot** happen - the
/// caller built it from the rank. And the coordinate came from a `Walk` over this very
/// shape, so the bounds cannot fail either. Two impossible checks, one `catch
/// unreachable`, on the hottest line in the library.
///
/// This computes the offset and returns it. The remaining safety is `data[offset]`,
/// which is where Zig puts it: a panic with a message in Debug and ReleaseSafe.
///
/// **Use it where the rank is an invariant you have already established** - after a
/// shape check, or inside a walk over the tensor's own shape. Use `flatIndex` where the
/// coordinate came from outside and might be wrong.
pub fn offsetOf(self: Self, indices: []const usize) usize {
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}
/// Read from a rank-2 tensor. **Does not return an error.**
///
/// WHY `at` HAS TO AND THIS DOES NOT
///
/// `at(indices: []const usize)` can fail for two reasons. One is a real bounds check.
/// The other is `indices.len != rank` - and that check exists ONLY because a slice hides
/// its arity from the type system. Every call site writes a literal `&.{ i, j }`; the
/// compiler could know it is two, and does not.
///
/// Naming the arity in the function removes that half entirely. What remains is the
/// bounds check, and **Zig already has a name for that**: it is what `data[idx]` does -
/// a panic with a message and a stack trace in Debug and ReleaseSafe, the same UB as any
/// other slice index in ReleaseFast.
///
/// Nobody writes `arr[i] catch unreachable`. This file had **105 `catch unreachable`**,
/// most of them on an `at` whose rank the author had checked three lines earlier, and
/// that noise is the single biggest thing that makes numeric Zig look worse than it is.
///
/// The safety is not weaker - it is Zig's, in the place Zig puts it.
pub fn at2(self: Self, i: usize, j: usize) T {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i)) +
self.strides[1] * @as(isize, @intCast(j));
return self.data[@intCast(offset)];
}
/// Write to a rank-2 tensor. The mirror of `at2`, and error-free for the same reason.
pub fn setAt2(self: Self, i: usize, j: usize, value: T) void {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i)) +
self.strides[1] * @as(isize, @intCast(j));
self.data[@intCast(offset)] = value;
}
/// Read from a rank-1 tensor, error-free.
pub fn at1(self: Self, i: usize) T {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i));
return self.data[@intCast(offset)];
}
/// Write to a rank-1 tensor, error-free.
pub fn setAt1(self: Self, i: usize, value: T) void {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i));
self.data[@intCast(offset)] = value;
}
/// Write `value` at `indices`.
pub fn setAt(self: Self, indices: []const usize, value: T) Error!void {
self.data[try self.flatIndex(indices)] = value;
}
/// Set every element to `value`.
pub fn fill(self: Self, value: T) void {
if (self.isContiguous()) {
@memset(self.data[self.base..][0..self.size()], value);
return;
}
var walker: [max_rank]usize = @splat(0);
var remaining: usize = self.size();
while (remaining > 0) : (remaining -= 1) {
self.data[self.offsetOf(walker[0..self.rank])] = value;
advance(self.shape[0..self.rank], walker[0..self.rank]);
}
}
/// The same storage seen as `shape`. Requires a contiguous tensor and an equal element
/// count; both are checked rather than assumed.
pub fn reshape(self: Self, shape: []const usize) Error!Self {
if (!self.isContiguous()) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
if (count != self.size()) {
return Error.ShapeMismatch;
}
// From `base`, not 0: reshaping a slice must keep addressing that slice.
return fromSlice(self.data[self.base..][0..self.size()], shape);
}
/// Whether two distinct index positions address the same element, which happens exactly
/// when a stretched axis (stride 0) has extent above 1.
///
/// An operation that WRITES must refuse an aliased destination. Many indices landing on
/// one element means the last write wins and every earlier one vanishes - a result that is
/// wrong in a way no bounds check catches. `map` and `zip` check this so a caller who
/// passes a broadcast view as an output gets an error rather than a plausible answer.
pub fn isAliased(self: Self) bool {
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
if (self.shape[axis] > 1 and self.strides[axis] == 0) {
return true;
}
}
return false;
}
/// A read-only view of this tensor stretched to `shape` by the broadcasting rule.
///
/// BROADCASTING IS A STRIDE OF ZERO, NOT A COPY
///
/// An axis whose extent is 1 is stretched by setting its stride to 0, so every position
/// along it addresses the same element. No storage is allocated and no data moves; a
/// `(3, 1, 5)` tensor viewed as `(3, 4, 5)` is the same 15 values read four times.
///
/// **WRITING THROUGH A STRETCHED VIEW IS ALMOST CERTAINLY A MISTAKE.** Many indices
/// alias one element, so the last write wins and the others vanish. This returns a view
/// because that is what the type system has; treat it as an input.
///
/// `isContiguous` reports false for any stretched axis with extent > 1, which is what
/// stops a linear walk from reading the underlying 15 values as if they were 60.
pub fn broadcastTo(self: Self, shape: []const usize) Error!Self {
if (shape.len > max_rank or shape.len < self.rank) {
return Error.ShapeMismatch;
}
var view: Self = .{
.data = self.data,
.shape = @splat(1),
.strides = @splat(0),
.rank = @intCast(shape.len),
// Stretching does not move where the view starts.
.base = self.base,
};
// Align at the TRAILING axis, as the broadcast rule does: an axis this tensor does not
// have behaves as extent 1, and so does an axis it has with extent 1. Both get a
// stride of 0.
for (shape, 0..) |extent, axis| {
const from_end: usize = shape.len - axis;
view.shape[axis] = extent;
if (self.rank < from_end) {
continue;
}
const own_axis: usize = self.rank - from_end;
const own_extent: usize = self.shape[own_axis];
if (own_extent == extent) {
view.strides[axis] = self.strides[own_axis];
} else if (own_extent != 1) {
return Error.ShapeMismatch;
}
}
return view;
}
/// A window into this tensor: `len` positions along `axis`, starting at `start`.
///
/// A SLICE IS AN OFFSET, WHICH IS WHY `base` EXISTS
///
/// Nothing is copied: the view keeps every stride and simply starts `start` steps further
/// into the data along one axis. A batch of 32 rows out of 4096 is a `Tensor` the same
/// size as any other, addressing 32 rows of the same buffer.
///
/// `start + len` must fit. Clamping instead would hand back a view SMALLER than asked
/// for, and the caller's next loop would read the shape it requested rather than the one
/// it got.
pub fn slice(self: Self, axis: usize, start: usize, len: usize) Error!Self {
if (axis >= self.rank) {
return Error.OutOfRange;
}
if (start + len > self.shape[axis]) {
return Error.OutOfRange;
}
var view: Self = self;
const step: isize = self.strides[axis] * @as(isize, @intCast(start));
const moved: isize = @as(isize, @intCast(self.base)) + step;
if (moved < 0) {
return Error.OutOfRange;
}
view.base = @intCast(moved);
view.shape[axis] = len;
return view;
}
/// A view with the axes reordered: `out` axis `i` is this tensor's axis `order[i]`.
///
/// EVERY `manip` OPERATION HERE IS METADATA
///
/// `permute`, `squeeze`, `unsqueeze`, `flatten` and `moveAxis` all return a view: they
/// permute or insert entries in two fixed-size arrays and copy nothing. That is why they
/// need no GPU kernel and no sweep row - **there is no arithmetic to disagree about.**
/// The cost appears later, when a permuted view has to be BOUND to a kernel and must be
/// densified first, which is the constraint recorded in section 10.3.
///
/// `order` must be a permutation of `0..rank`. A repeated axis would produce a view
/// where two indices alias, which is `isAliased`'s job to reject at write time - but
/// catching it here names the actual mistake.
pub fn permute(self: Self, order: []const usize) Error!Self {
if (order.len != self.rank) {
return Error.ShapeMismatch;
}
var seen: [max_rank]bool = @splat(false);
for (order) |axis| {
if (axis >= self.rank or seen[axis]) {
return Error.OutOfRange;
}
seen[axis] = true;
}
var view: Self = self;
for (order, 0..) |axis, i| {
view.shape[i] = self.shape[axis];
view.strides[i] = self.strides[axis];
}
return view;
}
/// A view with a size-1 axis removed.
///
/// Removing an axis whose extent is NOT 1 would change how many elements the view
/// addresses, so it is `ShapeMismatch` rather than a silent reshape.
pub fn squeeze(self: Self, axis: usize) Error!Self {
if (axis >= self.rank) {
return Error.OutOfRange;
}
if (self.shape[axis] != 1) {
return Error.ShapeMismatch;
}
var view: Self = self;
var i: usize = axis;
while (i + 1 < self.rank) : (i += 1) {
view.shape[i] = self.shape[i + 1];
view.strides[i] = self.strides[i + 1];
}
view.shape[self.rank - 1] = 1;
view.strides[self.rank - 1] = 0;
view.rank = @intCast(self.rank - 1);
return view;
}
/// A view with a size-1 axis inserted at `axis`. `axis == rank` appends.
///
/// The new axis gets a stride of 0, which is harmless because its extent is 1 - a
/// single position cannot step anywhere. `isAliased` only objects when extent exceeds 1.
pub fn unsqueeze(self: Self, axis: usize) Error!Self {
if (axis > self.rank or self.rank + 1 > max_rank) {
return Error.OutOfRange;
}
var view: Self = self;
var i: usize = self.rank;
while (i > axis) : (i -= 1) {
view.shape[i] = self.shape[i - 1];
view.strides[i] = self.strides[i - 1];
}
view.shape[axis] = 1;
view.strides[axis] = 0;
view.rank = @intCast(self.rank + 1);
return view;
}
/// A rank-1 view of the whole tensor. Contiguous only, for `reshape`'s reason.
pub fn flatten(self: Self) Error!Self {
return self.reshape(&.{self.size()});
}
/// A view with one axis moved from `from` to `to`, the others keeping their order.
pub fn moveAxis(self: Self, from: usize, to: usize) Error!Self {
if (from >= self.rank or to >= self.rank) {
return Error.OutOfRange;
}
var order: [max_rank]usize = @splat(0);
var k: usize = 0;
var i: usize = 0;
while (i < self.rank) : (i += 1) {
if (i == from) {
continue;
}
order[k] = i;
k += 1;
}
// Open a slot at `to` and drop the moved axis into it.
var j: usize = self.rank - 1;
while (j > to) : (j -= 1) {
order[j] = order[j - 1];
}
order[to] = from;
return self.permute(order[0..self.rank]);
}
/// A view with axes `a` and `b` exchanged. Shares storage; changes through one are visible
/// through the other.
pub fn transpose(self: Self, a: usize, b: usize) Error!Self {
if (a >= self.rank or b >= self.rank) {
return Error.OutOfRange;
}
var view: Self = self;
view.shape[a] = self.shape[b];
view.shape[b] = self.shape[a];
view.strides[a] = self.strides[b];
view.strides[b] = self.strides[a];
return view;
}
};
}
add/// `out = a + b`, broadcasting both to `out`'s shape.
pub fn add(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return x + y;
}
}.apply);
}
div is float-only: Zig has no single / for signed integers, and this
library does not choose between truncating, flooring and exact division on your behalf. Integer
division goes through zip.
scale and lerp take a runtime factor, which zip's
comptime function cannot close over, so they are written as loops.
lerp interpolates from the nearer endpoint:
const v: T = if (t <= 0.5) x + t * (y - x) else y - (1.0 - t) * (y - x);which returns exactly a at t = 0 and exactly b at
t = 1. Neither one-line form does both.
try zn.greater(f64, out, a, b); // 1 where a > b, else 0
try zn.less(f64, out, a, b);
try zn.equal(f64, out, a, b); // exact equality
try zn.notEqual(f64, out, a, b);
try zn.greaterEqual(f64, out, a, b);
try zn.lessEqual(f64, out, a, b);greater/// `out = 1 where a > b, else 0`. A MASK, not a bool tensor.
///
/// WHY COMPARISONS RETURN THE ELEMENT TYPE
///
/// A `Tensor(bool)` would be the obvious answer and is the wrong one here: a mask's whole purpose
/// is to be multiplied by something, and a bool tensor forces every consumer to convert. It also
/// cannot be a GPU buffer of the same element type, so the kernel would need a second dtype for
/// masks alone. Returning 0 or 1 in `T` keeps `mask * value` working on both backends with no
/// conversion anywhere.
pub fn greater(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
comptime requireNumeric(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return if (x > y) 1 else 0;
}
}.apply);
}
less/// `out = 1 where a < b, else 0`. See `greater`.
pub fn less(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
comptime requireNumeric(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return if (x < y) 1 else 0;
}
}.apply);
}
equal/// `out = 1 where a == b, else 0`.
///
/// EXACT equality, not a tolerance. A comparison with a built-in epsilon is a different
/// operation wearing this one's name, and a caller who wants tolerance can write
/// `less(|a - b|, eps)` and see the choice. NaN equals nothing, including itself.
pub fn equal(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
comptime requireNumeric(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return if (x == y) 1 else 0;
}
}.apply);
}
notEqual/// `out = 1 where a != b, else 0`. The complement of `equal`.
pub fn notEqual(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
) Error!void {
comptime requireNumeric(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return if (x != y) 1 else 0;
}
}.apply);
}
greaterEqual/// `out = 1 where a >= b, else 0`.
pub fn greaterEqual(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
) Error!void {
comptime requireNumeric(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return if (x >= y) 1 else 0;
}
}.apply);
}
lessEqual/// `out = 1 where a <= b, else 0`.
pub fn lessEqual(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
) Error!void {
comptime requireNumeric(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return if (x <= y) 1 else 0;
}
}.apply);
}
These return 1 or 0 in the element type, not a Tensor(bool). A mask
exists to be multiplied, and keeping it in T means mask * value works on
both backends with no conversion and no second buffer type.
The six masks partition in pairs — equal + notEqual,
greater + lessEqual, less + greaterEqual each sum to 1 at every element
— and for any pair of values exactly one of greater, less and
equal is 1. equal
compares exactly; a comparison with a built-in epsilon would be a different operation, and a caller
who wants tolerance writes less(|a - b|, eps) and sees the choice.
try zn.floor(f64, out, a); try zn.ceil(f64, out, a);
try zn.trunc(f64, out, a); try zn.round(f64, out, a);
try zn.abs(f64, out, a); try zn.neg(f64, out, a);
try zn.sign(f64, out, a); try zn.square(f64, out, a);floor/// `out = floor(a)`, elementwise.
pub fn floor(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return @floor(x);
}
}.apply);
}
ceil/// `out = ceil(a)`, elementwise.
pub fn ceil(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return @ceil(x);
}
}.apply);
}
trunc/// `out = scalarTrunc(a)` - toward zero, elementwise.
pub fn trunc(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
// `zm.trunc` gained a scalar branch this turn; using it keeps one implementation.
return scalarTrunc(x);
}
}.apply);
}
round/// `out = scalarRound(a)` - half away from zero, elementwise. Not banker's rounding: `zm.round` is
/// what the engine and the GPU both use, and matching it matters more than the tie rule.
pub fn round(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return scalarRound(x);
}
}.apply);
}
abs/// `out = |a|`, elementwise. `abs` for the same reason `exp` is not `abs`.
pub fn abs(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireNumeric(T);
return map(T, out, a, struct {
fn apply(x: T) T {
// `@abs` ON A SIGNED INTEGER RETURNS AN UNSIGNED ONE, which is correct of the
// builtin and wrong for a function whose output tensor has the input's type. The
// negation says the same thing and keeps the type.
//
// It also inherits the honest edge: the most negative integer has no positive
// counterpart, so `abs(minInt)` overflows rather than quietly returning itself.
if (comptime !isFloat(T)) {
return if (x < 0) -x else x;
}
return @abs(x);
}
}.apply);
}
neg/// `out = -a`, elementwise.
pub fn neg(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireNumeric(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return -x;
}
}.apply);
}
sign/// `out = -1, 0 or +1` by the sign of `a`. **Zero maps to zero**, not to +1: `sign` is used to
/// carry a direction, and giving a stationary point a direction is how a gradient acquires drift
/// that nothing accounts for.
pub fn sign(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireNumeric(T);
return map(T, out, a, struct {
fn apply(x: T) T {
if (x > 0) {
return 1;
}
if (x < 0) {
return -1;
}
return 0;
}
}.apply);
}
square/// `out = a * a`, elementwise. Not `pow(a, 2)`: a multiply is exact where a general power goes
/// through `exp(2*log(a))` and is neither exact nor defined for negatives.
pub fn square(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireNumeric(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return x * x;
}
}.apply);
}
trunc rounds toward zero and floor downward; they differ exactly on
the negatives. round is half away from zero. sign maps zero to zero,
so sign(x) * |x| == x holds for every input.
The suffixes exist because abs, exp, floor,
ceil, round, trunc, sqrt, log,
sin, cos, clamp and lerp are reserved math
words in this tree: a declaration may not shadow the zm function of that name.
try zn.exp(f64, out, a); try zn.log(f64, out, a);
try zn.expm1(f64, out, a); try zn.log1p(f64, out, a);
try zn.log2(f64, out, a); try zn.log10(f64, out, a);
try zn.sqrt(f64, out, a); try zn.cbrt(f64, out, a);
try zn.pow(f64, out, a, b); try zn.reciprocal(f64, out, a);
try zn.sin(f64, out, a); try zn.cos(f64, out, a);
try zn.atan2(f64, out, a, b);
try zn.hypot(f64, out, a, b);exp/// `out = e^a`, elementwise.
///
/// Named `exp`, not `exp`: `exp` is a reserved math word in this tree (it is `zm.exp`), and
/// the linter refuses a declaration that shadows one. znum's kernels use the same suffix for the
/// same reason, so the CPU and GPU names line up.
pub fn exp(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return @exp(x);
}
}.apply);
}
log/// `out = ln(a)`, elementwise. NaN below zero, -inf at zero. See `sqrt` on why neither is
/// clamped away.
pub fn log(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return @log(x);
}
}.apply);
}
expm1/// `out = e^a - 1`, elementwise.
///
/// Written as `@exp(x) - 1` rather than a dedicated series. A true `expm1` keeps precision
/// near zero, where `e^x - 1` cancels; this one does not, and says so rather than implying it.
/// It exists so that the same expression is written once, on both backends.
pub fn expm1(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return zm.expm1(x);
}
}.apply);
}
log1p/// `out = log(1 + a)`, elementwise. Same caveat as `expm1`: no extra precision near zero.
pub fn log1p(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return zm.log1p(x);
}
}.apply);
}
log2/// `out = log2(a)`, elementwise.
pub fn log2(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return scalarLog2(x);
}
}.apply);
}
log10/// `out = scalarLog10(a)`, elementwise.
pub fn log10(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return scalarLog10(x);
}
}.apply);
}
sqrt/// `out = sqrt(a)`, elementwise. **NaN for a negative input**, deliberately: the square root of a
/// negative real is undefined, and returning 0 or |a| would be inventing an answer the caller did
/// not ask for. Check the domain before calling if that matters.
pub fn sqrt(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return @sqrt(x);
}
}.apply);
}
cbrt/// `out = cube root of a`, elementwise, defined for negatives: `sign(x) * |x|^(1/3)`.
///
/// `pow(x, 1/3)` alone is NaN for a negative base. The real cube root of a negative number is
/// the negation of the cube root of its magnitude, and that is what this computes - matching
/// what every C library's `cbrt` returns.
pub fn cbrt(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return zm.cbrt(x);
}
}.apply);
}
pow/// `out = a ^ b` for a positive base, elementwise, as `exp(b * log a)`.
///
/// THE REAL-VALUED POWER, NOT THE C LIBRARY'S. `zm.pow` is `f32`-only and carries the C
/// special cases - a negative base with an integer exponent, `pow(x, 0)` for any `x`. This is the
/// definition that is the same expression on every backend and every float width: NaN for a
/// negative base, whatever the exponent, and `0^y` handled explicitly because `log 0` is `-inf`.
/// Delegates to `zm.pow`, so a negative base with a whole exponent works: `pow(-2, 3)` is -8.
pub fn pow(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
) Error!void {
comptime requireFloat(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
// `zm.pow`, NOT `exp(y * log(x))` - THEY DISAGREE ON NEGATIVE BASES
//
// The exponential form needs `log(x)`, which is NaN for x < 0. So `pow(-2, 3)` came
// back NaN where the answer is -8, and `pow(-2, 2)` NaN where it is 4. The doc used
// to call that deliberate; it was not defensible once `zm.pow` existed and returned
// both correctly.
//
// **The correspondence gate pairs this with `zm.pow` by NAME and never compared a
// value**, so the two could drift apart indefinitely. The test now checks agreement
// on a spread that includes negative bases, which is what the pairing was supposed
// to mean.
return zm.pow(x, y);
}
}.apply);
}
reciprocal/// `out = 1 / a`, elementwise. Infinite at zero on both backends, which the sweep's comparison
/// treats as agreement when the signs match.
pub fn reciprocal(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return 1.0 / x;
}
}.apply);
}
hypot/// `out = sqrt(a*a + b*b)`, elementwise.
pub fn hypot(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
) Error!void {
comptime requireFloat(T);
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return @sqrt(x * x + y * y);
}
}.apply);
}
pow is exp(b · log a) — the real-valued power, NaN for a
negative base whatever the exponent, and the same expression on every backend and float width.
cbrt is defined for negatives as sign(x) · |x|1/3,
which pow(x, 1/3) is not; this is exact for cbrt(x)³ == x on both
signs. expm1 and log1p are written as exp(x) − 1 and
log(1 + x) and carry no extra precision near zero.
sqrt and log return NaN outside their domain rather than clamping,
and log returns −inf at zero. reciprocal is infinite at zero.
Check the domain before calling if that matters.
atan2 takes (y, x), matching every C library.
try zn.relu(f32, out, a);
try zn.sigmoid(f32, out, a);
try zn.tanh(f32, out, a);
try zn.gelu(f32, out, a); // tanh approximation
try zn.softplus(f32, out, a);
try zn.silu(f32, out, a); // x * sigmoid(x)
try zn.leakyRelu(f32, out, a, 0.01);
try zn.elu(f32, out, a, 1.0);relu/// `out = max(a, 0)`, elementwise.
pub fn relu(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return if (x > 0) x else 0;
}
}.apply);
}
sigmoid/// `out = 1 / (1 + e^-a)`, elementwise, through `zm.sigmoid`.
///
/// THE HOST AND THE GPU RUN THE SAME SOURCE, NOT THE SAME INSTRUCTIONS. `zm` compiles a
/// different branch for SPIR-V (`@exp`) than for the host (`std.math.exp`), so the two agree to
/// within a rounding step rather than exactly. Measured on an Adreno phone over 4096 f32:
/// `relu` 0, `sigmoid` and `tanh` 1.19e-7 (one ULP), `gelu` 2.38e-7 (two - it is built on
/// `tanh`). A caller needing bit-identical cross-backend results does not have it here; a caller
/// needing a CPU reference that predicts the GPU to a rounding step does.
pub fn sigmoid(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return zm.sigmoid(x);
}
}.apply);
}
tanh/// `out = tanh(a)`, elementwise.
pub fn tanh(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return zm.tanh(x);
}
}.apply);
}
gelu/// `out = gelu(a)`, elementwise - the tanh approximation PyTorch spells
/// `gelu(approximate="tanh")`.
pub fn gelu(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return zm.gelu(x);
}
}.apply);
}
softplus/// `out = log(1 + e^a)`, elementwise - a smooth relu.
///
/// GUARDED ABOVE 20. `e^x` overflows f32 near 88 and `log(1 + inf)` is inf, but softplus(x)
/// is within one f32 ULP of `x` itself well before that. Returning `x` above the threshold is
/// both correct and the only way to stay finite, and the CPU and the kernel must guard at the
/// SAME point or the two disagree by an infinity on one element.
pub fn softplus(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
// SHIFTED, AND WITH `log1p` - WHICH BEATS BOTH EARLIER FORMS
//
// The identity: for any x,
//
// log(1 + e^x) = max(x, 0) + log(1 + e^-|x|)
//
// because for x > 0, log(1 + e^x) = log(e^x (e^-x + 1)) = x + log(1 + e^-x), and for
// x <= 0 the max is zero and |x| = -x leaves the original. So the exponential is only
// ever taken of a NEGATIVE number and cannot overflow.
//
// Two earlier forms and what each lost, measured at f64:
//
// `if (x > 20) x else log(1 + e^x)` - ours until now. A DISCONTINUITY at 20:
// relative error 1.0e-10 just above it, where f64 holds 2.2e-16.
//
// `max(x,0) + log(1 + e^-|x|)` - znum's. Right at the positive end, and still
// **exactly zero from x = -50 down**, because `1 + 1.9e-22` rounds to 1 and its
// logarithm is 0. The true value there is 1.9e-22.
//
// `log1p` takes the small argument directly and never forms the 1, so the tail
// survives to the smallest subnormal. It is the same trick `logSumExp` uses one
// screen away, and the same one `log(softmax(x))` needed.
const shifted: T = @exp(-@abs(x));
return @max(x, 0) + zm.log1p(shifted);
}
}.apply);
}
silu/// `out = a * sigmoid(a)`, elementwise - SiLU, also called swish.
pub fn silu(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return x * zm.sigmoid(x);
}
}.apply);
}
leakyRelu/// `out = a` where positive, `slope * a` elsewhere.
///
/// The slope is a parameter with no default. PyTorch uses 0.01, but a leaky relu whose slope
/// the caller did not choose is a hyperparameter chosen by whoever wrote the library.
pub fn leakyRelu(comptime T: type, out: Tensor(T), a: Tensor(T), slope: T) Error!void {
comptime requireFloat(T);
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const va: Tensor(T) = try a.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const x: T = va.data[va.offsetOf(at)];
out.data[out.offsetOf(at)] = if (x > 0) x else slope * x;
}
}
elu/// `out = a` where positive, `alpha * (e^a - 1)` elsewhere.
///
/// Continuous at zero for every alpha - both branches give 0 - which is what makes it usable as
/// an activation at all. Its DERIVATIVE is continuous only at alpha = 1, and that is a property
/// of the function rather than something this implementation should hide.
pub fn elu(comptime T: type, out: Tensor(T), a: Tensor(T), alpha: T) Error!void {
comptime requireFloat(T);
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const va: Tensor(T) = try a.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const x: T = va.data[va.offsetOf(at)];
out.data[out.offsetOf(at)] =
if (x > 0) x else alpha * (@exp(x) - 1.0);
}
}
sigmoid, tanh and gelu call the corresponding
zm function, so the CPU and any GPU kernel run the same source.
softplus returns x above 20: e^x overflows
f32 near 88, and softplus is within one ULP of x well before that. The
GPU kernel guards at the same threshold.
leakyRelu and elu take their parameter with no default. Both
branches meet at zero for any parameter value.
try zn.reluGrad(f32, out, x, dy); // x is the forward INPUT
try zn.sigmoidGrad(f32, out, y, dy); // y is the forward OUTPUT
try zn.tanhGrad(f32, out, y, dy); // y is the forward OUTPUTreluGrad/// `out = if (x > 0) grad else 0` - the gradient of `relu` with respect to its input.
///
/// It takes the FORWARD INPUT, not the forward output. Both work for `relu` because the output
/// is positive exactly where the input was, but only the input generalises: `leakyRelu` and
/// `elu` have outputs that do not determine their own gradient. Taking the input here means the
/// signature does not have to change when they arrive.
pub fn reluGrad(
comptime T: type,
out: Tensor(T),
x: Tensor(T),
grad: Tensor(T),
) Error!void {
comptime requireFloat(T);
return zip(T, out, x, grad, struct {
fn apply(xv: T, gv: T) T {
return if (xv > 0) gv else 0;
}
}.apply);
}
sigmoidGrad/// `out = grad * y * (1 - y)` - sigmoid's gradient, from its OUTPUT `y`.
///
/// TAKES THE OUTPUT, where `reluGrad` takes the input. That asymmetry is not an oversight: a
/// sigmoid's derivative is `y(1-y)` in terms of its own output, so passing `y` costs one multiply
/// where passing `x` would mean recomputing the sigmoid. `relu`'s derivative is not expressible
/// in its output for the family it belongs to, which is why that one takes `x`. The rule is
/// "whichever the derivative is actually a function of", and it differs per activation.
pub fn sigmoidGrad(
comptime T: type,
out: Tensor(T),
y: Tensor(T),
grad: Tensor(T),
) Error!void {
comptime requireFloat(T);
return zip(T, out, y, grad, struct {
fn apply(yv: T, gv: T) T {
return gv * yv * (1.0 - yv);
}
}.apply);
}
tanhGrad/// `out = grad * (1 - y*y)` - tanh's gradient, also from its output. See `sigmoidGrad`.
pub fn tanhGrad(
comptime T: type,
out: Tensor(T),
y: Tensor(T),
grad: Tensor(T),
) Error!void {
comptime requireFloat(T);
return zip(T, out, y, grad, struct {
fn apply(yv: T, gv: T) T {
return gv * (1.0 - yv * yv);
}
}.apply);
}
The argument differs by activation because the derivative does. A sigmoid's derivative is
y(1-y) in terms of its own output, so passing y is one multiply where
passing x would mean recomputing the sigmoid. A relu's derivative is a test on its
input, and taking the input generalises to leakyRelu and elu, whose
outputs do not determine their gradients.
All three are checked against a central finite difference of the forward function, in
f64.
try zn.logSoftmaxRows(f32, out, logits);logSoftmaxRows/// The logarithm of the softmax, computed without ever forming the softmax.
///
/// `log(softmax(x))` IS NOT `log` OF `softmax(x)`, NUMERICALLY
///
/// Taking the softmax and then the logarithm loses the small probabilities entirely: a value
/// twenty below the row maximum softmaxes to `2e-9`, which survives, and one ninety below
/// underflows to **exactly zero** whose logarithm is negative infinity. The true answer is -90,
/// a perfectly ordinary number.
///
/// The identity avoids it: `log(softmax(x))_i = x_i - m - log(sum(exp(x - m)))`, which is the
/// row's log-sum-exp subtracted from each element. Nothing is ever exponentiated without the
/// maximum removed, and nothing small is ever stored.
///
/// This is what a cross-entropy should be built on, and why every framework has it as a separate
/// function rather than a composition. znum has `logSoftmax` too; the difference here is that
/// `logSumExpAxis` already existed, so this is four lines rather than a second derivation.
pub fn logSoftmaxRows(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
if (out.shape[0] != a.shape[0] or out.shape[1] != a.shape[1]) {
return Error.ShapeMismatch;
}
for (0..a.shape[0]) |row| {
var largest: T = zm.lowest(T);
for (0..a.shape[1]) |col| {
largest = @max(largest, try a.at(&.{ row, col }));
}
var total: CompensatedSum(T) = .{};
for (0..a.shape[1]) |col| {
total.add(@exp((try a.at(&.{ row, col })) - largest));
}
const offset: T = largest + @log(total.value());
for (0..a.shape[1]) |col| {
try out.setAt(&.{ row, col }, (try a.at(&.{ row, col })) - offset);
}
}
}
log(softmax(x)) is not log of softmax(x).
exp(-x) in f32 reaches zero near x = 103.3, where the smallest
subnormal is. A value 110 below the row maximum softmaxes to exactly zero, and the
logarithm of zero is negative infinity. The true answer is −110, a perfectly ordinary
number.
The composition does not lose precision — it loses the value. The identity avoids it
entirely: log(softmax(x))ᵢ = xᵢ − logSumExp(x), so nothing is exponentiated
without the maximum removed and nothing small is ever stored.
The test asserts four things: that the composition really does produce -inf at
−110; that this gives a finite −110; that the two agree where the composition
works, so the identity costs nothing; and that exponentiating the result sums to one, which is the
definition and the strongest check available without a second implementation.
One number in that test was wrong first. one might write −90 from memory and it failed, correctly
— exp(-90) is 8.2e−40, a subnormal that survives. The threshold is
a property of the format, not a round number, and guessing it is the same mistake as
guessing a tolerance.
try zn.compareScalar(i32, out, a, 3, .greater_or_equal);compareScalar/// Every element compared against one value, with the comparison named.
///
/// ONE FUNCTION WHERE znum HAS SIX
///
/// znum has `eqScalar`, `neScalar`, `ltScalar`, `leScalar`, `gtScalar` and `geScalar`. They are
/// the same loop with a different operator, and six names means six places to look, six docs to
/// keep true, and six chances for one of them to disagree with the others about an edge.
///
/// A named comparison says the same thing at the call site - `compareScalar(out, a, 0, .greater)`
/// reads as well as `gtScalar(a, 0)` - and there is one implementation to be right.
pub fn compareScalar(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
value: T,
how: Comparison,
) Error!void {
comptime requireNumeric(T);
if (out.rank != a.rank) {
return Error.ShapeMismatch;
}
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
const x: T = try a.at(at);
const holds: bool = switch (how) {
.equal => x == value,
.not_equal => x != value,
.less => x < value,
.less_or_equal => x <= value,
.greater => x > value,
.greater_or_equal => x >= value,
};
try out.setAt(at, if (holds) 1 else 0);
}
}
znum has eqScalar, neScalar, ltScalar,
leScalar, gtScalar and geScalar. They are the same loop with
a different operator — six names, six docs to keep true, and six chances for one to disagree
with the others about the edge where the value is equal. A named comparison reads as well at the
call site and there is one implementation to be right about.
bitwiseAnd(2, 1) == 0 // they share no bits
logicalAnd(2, 1) == 1 // both are trueNot an exotic edge — two small positive integers. The same trap as the two
nots beside them, and the reason both families are named in full.
const total: f64 = zn.sumAll(f64, t);
const fast: f64 = zn.sumAllFast(f64, t);
const avg: f64 = try zn.meanAll(f64, t);
const lo: f64 = try zn.minAll(f64, t);
const hi: f64 = try zn.maxAll(f64, t);sumAll/// The sum of every element, in compensated arithmetic.
///
/// WHY THIS IS NOT A `for` LOOP WITH A RUNNING TOTAL
///
/// A naive sum loses the low bits of every addend once the running total grows past them. Summing
/// a million values near 1.0 in `f32`, the total reaches ~10^6 while each addend is ~1, and the
/// last additions round away most of their own magnitude: the error grows with n, and the answer
/// is quietly wrong in the direction of "too small".
///
/// Neumaier's variant of Kahan summation carries the lost part in a second accumulator and adds
/// it back at the end. It costs three extra flops per element - nothing next to the memory
/// traffic - and it handles the case plain Kahan gets wrong, where the ADDEND is larger than the
/// running total. `zn.sumAll` over a million ones in `f32` is exact; the naive loop is not.
pub fn sumAll(comptime T: type, a: Tensor(T)) T {
comptime requireNumeric(T);
var total: CompensatedSum(T) = .{};
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
total.add(a.data[a.offsetOf(at)]);
}
return total.value();
}
sumAllFast/// The sum of every element by PAIRWISE recursion - the fast variant.
///
/// 3.3x FASTER THAN `sumAll`, WITH A CLIFF
///
/// Measured on 168M `f32`, one core: **pairwise 1 398 M/s, Neumaier 425 M/s**, and a plain
/// vectorised `+=` loop 704 M/s in between. Pairwise beats even the naive loop because the
/// recursion produces independent accumulation chains the CPU pipelines.
///
/// Its error grows as O(log n) rather than O(n), which is enough for almost everything - a
/// million ones in `f32` comes back exactly.
///
/// **AND THEN IT FALLS OFF A CLIFF UNDER CANCELLATION.** Summing
/// `[1e8, 1x100, -1e8]`, whose exact answer is 100, measured across thresholds:
///
/// threshold 128 -> 0 threshold 8 -> 104
/// threshold 32 -> 72 threshold 2 -> 96
///
/// Not merely inaccurate - **erratic**, and at the usual threshold completely wrong. The two
/// large values land in different halves and each swamps the small values beside it. No threshold
/// fixes it, because the failure is in the shape of the recursion rather than its depth.
///
/// So this is the named variant and `sumAll` is the default. Reach for it when the data is
/// known not to cancel - magnitudes, squares, counts - and the array is large enough for 3.3x to
/// matter.
pub fn sumAllFast(comptime T: type, a: Tensor(T)) T {
comptime requireNumeric(T);
if (!a.isContiguous()) {
// A strided view has no contiguous halves to recurse on; fall back to the exact path
// rather than silently walking the wrong elements.
return sumAll(T, a);
}
return pairwiseSum(T, a.data[a.base..][0..a.size()]);
}
meanAll/// The arithmetic mean of every element, or `DomainError` for an empty tensor.
///
/// An empty mean is undefined rather than zero or NaN, and saying so is the caller's cue to
/// decide what an empty batch means - which is a question only the caller can answer.
pub fn meanAll(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n == 0) {
return Error.DomainError;
}
const divisor: T = @floatFromInt(n);
return sumAll(T, a) / divisor;
}
minAll/// The smallest element, or `DomainError` for an empty tensor.
///
/// NaN is not ordered, so a NaN in the data makes every comparison false and the result is
/// whichever element happened to come first. Callers who care must check their inputs; this
/// function will not invent an ordering the type does not have.
pub fn minAll(comptime T: type, a: Tensor(T)) Error!T {
return extremum(T, a, true);
}
maxAll/// The largest element, or `DomainError` for an empty tensor. See `minAll` on NaN.
pub fn maxAll(comptime T: type, a: Tensor(T)) Error!T {
return extremum(T, a, false);
}
All work on any view, strided or stretched. meanAll, minAll and
maxAll return DomainError for an empty tensor. NaN is not ordered, so
minAll and maxAll over data containing one return whichever element came
first.
sumAll carries the part each addition loses in a second accumulator:
const t: T = total + x;
if (@abs(total) >= @abs(x)) {
compensation += (total - t) + x;
} else {
compensation += (x - t) + total;
}
total = t;
// ... and the answer is `total + compensation`This is Neumaier's variant of Kahan summation; the branch handles the case where the addend is
larger than the running total, which plain Kahan gets wrong. It costs three extra operations per
element. The loop is CompensatedSum(T), and every accumulation in the library — the
reductions, matmul's inner product, the statistics, every loss — goes through
it:
var total: zn.CompensatedSum(f64) = .{};
for (values) |x| total.add(x);
const sum = total.value();CompensatedSum/// A running sum that carries the rounding error of every addition in a second accumulator.
///
/// ONE IMPLEMENTATION OF NEUMAIER SUMMATION, USED BY EVERY ACCUMULATION IN THIS FILE
///
/// The eight lines below had been copied into ten places - `sumAll`, `matmul`'s inner product,
/// `layerNormRows` twice, `variance`, `covariance`, `dotAll`, `mseLoss`, `rootMeanSquare` - and the five
/// losses written later had none of them, because copying eight lines is exactly the kind of
/// thing that gets skipped. A type with an `add` is what makes compensation the default rather
/// than a discipline.
///
/// Neumaier rather than Kahan: the branch handles an addend larger than the running total,
/// which plain Kahan gets wrong. The cost is three extra floating-point operations per element.
pub fn CompensatedSum(comptime T: type) type {
comptime requireNumeric(T);
return struct {
total: T = 0,
compensation: T = 0,
/// THE COMPENSATION IS A NO-OP FOR INTEGERS, AND THAT IS WHY IT LIVES HERE
///
/// Integer addition is EXACT - there is no rounding, so there is nothing to compensate
/// and the correction term would be zero at every step. Branching here rather than at
/// every call site means each of the dozen reductions built on this type widens to
/// integers for free, instead of a dozen copies of the same comptime `if`.
///
/// The same lesson as `Walk`: **a type with a method makes the right thing the default**.
pub fn add(self: *@This(), x: T) void {
if (comptime !isFloat(T)) {
self.total += x;
return;
}
const t: T = self.total + x;
if (@abs(self.total) >= @abs(x)) {
self.compensation += (self.total - t) + x;
} else {
self.compensation += (x - t) + self.total;
}
self.total = t;
}
pub fn value(self: @This()) T {
return self.total + self.compensation;
}
};
}
sumAllFast instead recurses, halving until 128 elements and then accumulating
plainly. Measured on 168M f32 values on one core:
| Method | Throughput | [1e8, 1×100, -1e8], exact answer 100 |
|---|---|---|
sumAll (Neumaier) | 425 M/s | 100 |
plain += loop | 704 M/s | 0 |
sumAllFast (pairwise) | 1,398 M/s | 0 at threshold 128; 72, 104, 96 at 32, 8, 2 |
Pairwise error grows as O(log n) rather than O(n) and a million ones in f32
comes back exactly from either. The difference appears under cancellation, where pairwise is not
merely less accurate but erratic. Use sumAllFast when the data cannot cancel:
magnitudes, squares, counts.
matmul's inner product and layerNormRows's two accumulations use the
compensated form for the same reason.
const cols = try zn.Tensor(f64).alloc(a, &.{3}); // (2,3) summed down axis 0
try zn.sumAxis(f64, cols, t, 0);Tensor/// A view of `T`-typed storage as a multidimensional array.
///
/// A TENSOR NEVER OWNS ITS STORAGE, AND THAT IS THE DESIGN
///
/// There is no ownership flag and no `deinit`. `data` is a slice the tensor addresses; whoever
/// allocated it frees it, with the allocator that allocated it. Usually that is an arena, in
/// which case nobody frees anything and `arena.deinit()` ends the matter.
///
/// The alternative - an `owns: bool` and a `deinit` that is a no-op for views - is the shape
/// this library was warned about before it had a tensor: both of its failure modes are SILENT.
/// `deinit` on a view does nothing and reads as if it did something; a missing `deinit` on an
/// owner just grows memory. Having no `deinit` at all means a free appears exactly where an
/// allocation appears, and a view has neither.
///
/// Strides are in ELEMENTS and signed, so a reversed or negatively-strided view is expressible
/// without changing this type. `at` computes a signed offset and casts once.
pub fn Tensor(comptime T: type) type {
return struct {
const Self = @This();
/// The storage this tensor addresses. Not owned - see the type's doc comment.
data: []T,
/// Extent of each axis. Only the first `rank` entries are meaningful.
shape: [max_rank]usize,
/// Step in ELEMENTS between consecutive positions along each axis.
strides: [max_rank]isize,
rank: u8,
/// WHERE THIS VIEW STARTS INSIDE `data`
///
/// Every view until now began at element 0, which is why `slice` could not exist: a
/// window into the middle of a tensor IS an offset. With this, `slice` and the rest of
/// znum's `index` namespace become views rather than copies.
///
/// IT ALSO WEAKENS WHAT `data.len` MEANS. `data` is the whole backing store; `base`
/// and the strides say which of it this view addresses. Everything that walked
/// `data[0..size()]` assuming those coincide had to be corrected - five sites, all of
/// them fast paths that skip the index walk, which is exactly where such an assumption
/// hides.
base: usize = 0,
/// A dense row-major tensor over `shape`, allocated from `gpa`.
///
/// The caller frees `result.data` with the same allocator - or allocates from an arena
/// and frees nothing.
pub fn alloc(gpa: Allocator, shape: []const usize) Error!Self {
if (shape.len > max_rank) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
const storage: []T = gpa.alloc(T, count) catch return Error.OutOfMemory;
return fromSlice(storage, shape);
}
/// A dense row-major view of `data` as `shape`. `data.len` must be exactly the product of
/// the extents - a slice that is merely LARGE ENOUGH is rejected, because the usual reason
/// for the mismatch is a shape that is wrong rather than storage that is generous.
pub fn fromSlice(data: []T, shape: []const usize) Error!Self {
if (shape.len > max_rank) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
if (count != data.len) {
return Error.ShapeMismatch;
}
var self: Self = .{
.data = data,
.shape = @splat(1),
.strides = @splat(0),
.rank = @intCast(shape.len),
.base = 0,
};
// Row-major: the last axis is contiguous, and each earlier stride is the product of
// the extents to its right.
var stride: isize = 1;
var axis: usize = shape.len;
while (axis > 0) {
axis -= 1;
self.shape[axis] = shape[axis];
self.strides[axis] = stride;
stride *= @intCast(shape[axis]);
}
return self;
}
/// The number of elements this tensor addresses.
pub fn size(self: Self) usize {
var count: usize = 1;
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
count *= self.shape[axis];
}
return count;
}
/// Whether the strides are dense row-major, so `data[0..size()]` is the tensor in order.
///
/// An operation that walks memory linearly must ask this rather than assume it: a view
/// produced by `transpose` addresses the same storage in a different order, and a linear
/// walk over it silently reads the wrong elements rather than failing.
pub fn isContiguous(self: Self) bool {
var expected: isize = 1;
var axis: usize = self.rank;
while (axis > 0) {
axis -= 1;
if (self.shape[axis] != 1 and self.strides[axis] != expected) {
return false;
}
expected *= @intCast(self.shape[axis]);
}
return true;
}
/// The offset into `data` of the element at `indices`.
pub fn flatIndex(self: Self, indices: []const usize) Error!usize {
if (indices.len != self.rank) {
return Error.ShapeMismatch;
}
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
if (i >= self.shape[axis]) {
return Error.OutOfRange;
}
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}
/// The element at `indices`.
pub fn at(self: Self, indices: []const usize) Error!T {
return self.data[try self.flatIndex(indices)];
}
/// The flat offset for a coordinate, with no error.
///
/// THE 45 `flatIndex(...) catch unreachable` SITES ARE ALL THIS SHAPE
///
/// const at: []const usize = walker[0..out.rank];
/// out.data[out.offsetOf(at)] = ...
///
/// The slice is `walker[0..rank]`, so `indices.len != rank` **cannot** happen - the
/// caller built it from the rank. And the coordinate came from a `Walk` over this very
/// shape, so the bounds cannot fail either. Two impossible checks, one `catch
/// unreachable`, on the hottest line in the library.
///
/// This computes the offset and returns it. The remaining safety is `data[offset]`,
/// which is where Zig puts it: a panic with a message in Debug and ReleaseSafe.
///
/// **Use it where the rank is an invariant you have already established** - after a
/// shape check, or inside a walk over the tensor's own shape. Use `flatIndex` where the
/// coordinate came from outside and might be wrong.
pub fn offsetOf(self: Self, indices: []const usize) usize {
var offset: isize = @intCast(self.base);
for (indices, 0..) |i, axis| {
offset += self.strides[axis] * @as(isize, @intCast(i));
}
return @intCast(offset);
}
/// Read from a rank-2 tensor. **Does not return an error.**
///
/// WHY `at` HAS TO AND THIS DOES NOT
///
/// `at(indices: []const usize)` can fail for two reasons. One is a real bounds check.
/// The other is `indices.len != rank` - and that check exists ONLY because a slice hides
/// its arity from the type system. Every call site writes a literal `&.{ i, j }`; the
/// compiler could know it is two, and does not.
///
/// Naming the arity in the function removes that half entirely. What remains is the
/// bounds check, and **Zig already has a name for that**: it is what `data[idx]` does -
/// a panic with a message and a stack trace in Debug and ReleaseSafe, the same UB as any
/// other slice index in ReleaseFast.
///
/// Nobody writes `arr[i] catch unreachable`. This file had **105 `catch unreachable`**,
/// most of them on an `at` whose rank the author had checked three lines earlier, and
/// that noise is the single biggest thing that makes numeric Zig look worse than it is.
///
/// The safety is not weaker - it is Zig's, in the place Zig puts it.
pub fn at2(self: Self, i: usize, j: usize) T {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i)) +
self.strides[1] * @as(isize, @intCast(j));
return self.data[@intCast(offset)];
}
/// Write to a rank-2 tensor. The mirror of `at2`, and error-free for the same reason.
pub fn setAt2(self: Self, i: usize, j: usize, value: T) void {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i)) +
self.strides[1] * @as(isize, @intCast(j));
self.data[@intCast(offset)] = value;
}
/// Read from a rank-1 tensor, error-free.
pub fn at1(self: Self, i: usize) T {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i));
return self.data[@intCast(offset)];
}
/// Write to a rank-1 tensor, error-free.
pub fn setAt1(self: Self, i: usize, value: T) void {
const offset: isize = @as(isize, @intCast(self.base)) +
self.strides[0] * @as(isize, @intCast(i));
self.data[@intCast(offset)] = value;
}
/// Write `value` at `indices`.
pub fn setAt(self: Self, indices: []const usize, value: T) Error!void {
self.data[try self.flatIndex(indices)] = value;
}
/// Set every element to `value`.
pub fn fill(self: Self, value: T) void {
if (self.isContiguous()) {
@memset(self.data[self.base..][0..self.size()], value);
return;
}
var walker: [max_rank]usize = @splat(0);
var remaining: usize = self.size();
while (remaining > 0) : (remaining -= 1) {
self.data[self.offsetOf(walker[0..self.rank])] = value;
advance(self.shape[0..self.rank], walker[0..self.rank]);
}
}
/// The same storage seen as `shape`. Requires a contiguous tensor and an equal element
/// count; both are checked rather than assumed.
pub fn reshape(self: Self, shape: []const usize) Error!Self {
if (!self.isContiguous()) {
return Error.UnsupportedShape;
}
var count: usize = 1;
for (shape) |extent| {
count *= extent;
}
if (count != self.size()) {
return Error.ShapeMismatch;
}
// From `base`, not 0: reshaping a slice must keep addressing that slice.
return fromSlice(self.data[self.base..][0..self.size()], shape);
}
/// Whether two distinct index positions address the same element, which happens exactly
/// when a stretched axis (stride 0) has extent above 1.
///
/// An operation that WRITES must refuse an aliased destination. Many indices landing on
/// one element means the last write wins and every earlier one vanishes - a result that is
/// wrong in a way no bounds check catches. `map` and `zip` check this so a caller who
/// passes a broadcast view as an output gets an error rather than a plausible answer.
pub fn isAliased(self: Self) bool {
var axis: usize = 0;
while (axis < self.rank) : (axis += 1) {
if (self.shape[axis] > 1 and self.strides[axis] == 0) {
return true;
}
}
return false;
}
/// A read-only view of this tensor stretched to `shape` by the broadcasting rule.
///
/// BROADCASTING IS A STRIDE OF ZERO, NOT A COPY
///
/// An axis whose extent is 1 is stretched by setting its stride to 0, so every position
/// along it addresses the same element. No storage is allocated and no data moves; a
/// `(3, 1, 5)` tensor viewed as `(3, 4, 5)` is the same 15 values read four times.
///
/// **WRITING THROUGH A STRETCHED VIEW IS ALMOST CERTAINLY A MISTAKE.** Many indices
/// alias one element, so the last write wins and the others vanish. This returns a view
/// because that is what the type system has; treat it as an input.
///
/// `isContiguous` reports false for any stretched axis with extent > 1, which is what
/// stops a linear walk from reading the underlying 15 values as if they were 60.
pub fn broadcastTo(self: Self, shape: []const usize) Error!Self {
if (shape.len > max_rank or shape.len < self.rank) {
return Error.ShapeMismatch;
}
var view: Self = .{
.data = self.data,
.shape = @splat(1),
.strides = @splat(0),
.rank = @intCast(shape.len),
// Stretching does not move where the view starts.
.base = self.base,
};
// Align at the TRAILING axis, as the broadcast rule does: an axis this tensor does not
// have behaves as extent 1, and so does an axis it has with extent 1. Both get a
// stride of 0.
for (shape, 0..) |extent, axis| {
const from_end: usize = shape.len - axis;
view.shape[axis] = extent;
if (self.rank < from_end) {
continue;
}
const own_axis: usize = self.rank - from_end;
const own_extent: usize = self.shape[own_axis];
if (own_extent == extent) {
view.strides[axis] = self.strides[own_axis];
} else if (own_extent != 1) {
return Error.ShapeMismatch;
}
}
return view;
}
/// A window into this tensor: `len` positions along `axis`, starting at `start`.
///
/// A SLICE IS AN OFFSET, WHICH IS WHY `base` EXISTS
///
/// Nothing is copied: the view keeps every stride and simply starts `start` steps further
/// into the data along one axis. A batch of 32 rows out of 4096 is a `Tensor` the same
/// size as any other, addressing 32 rows of the same buffer.
///
/// `start + len` must fit. Clamping instead would hand back a view SMALLER than asked
/// for, and the caller's next loop would read the shape it requested rather than the one
/// it got.
pub fn slice(self: Self, axis: usize, start: usize, len: usize) Error!Self {
if (axis >= self.rank) {
return Error.OutOfRange;
}
if (start + len > self.shape[axis]) {
return Error.OutOfRange;
}
var view: Self = self;
const step: isize = self.strides[axis] * @as(isize, @intCast(start));
const moved: isize = @as(isize, @intCast(self.base)) + step;
if (moved < 0) {
return Error.OutOfRange;
}
view.base = @intCast(moved);
view.shape[axis] = len;
return view;
}
/// A view with the axes reordered: `out` axis `i` is this tensor's axis `order[i]`.
///
/// EVERY `manip` OPERATION HERE IS METADATA
///
/// `permute`, `squeeze`, `unsqueeze`, `flatten` and `moveAxis` all return a view: they
/// permute or insert entries in two fixed-size arrays and copy nothing. That is why they
/// need no GPU kernel and no sweep row - **there is no arithmetic to disagree about.**
/// The cost appears later, when a permuted view has to be BOUND to a kernel and must be
/// densified first, which is the constraint recorded in section 10.3.
///
/// `order` must be a permutation of `0..rank`. A repeated axis would produce a view
/// where two indices alias, which is `isAliased`'s job to reject at write time - but
/// catching it here names the actual mistake.
pub fn permute(self: Self, order: []const usize) Error!Self {
if (order.len != self.rank) {
return Error.ShapeMismatch;
}
var seen: [max_rank]bool = @splat(false);
for (order) |axis| {
if (axis >= self.rank or seen[axis]) {
return Error.OutOfRange;
}
seen[axis] = true;
}
var view: Self = self;
for (order, 0..) |axis, i| {
view.shape[i] = self.shape[axis];
view.strides[i] = self.strides[axis];
}
return view;
}
/// A view with a size-1 axis removed.
///
/// Removing an axis whose extent is NOT 1 would change how many elements the view
/// addresses, so it is `ShapeMismatch` rather than a silent reshape.
pub fn squeeze(self: Self, axis: usize) Error!Self {
if (axis >= self.rank) {
return Error.OutOfRange;
}
if (self.shape[axis] != 1) {
return Error.ShapeMismatch;
}
var view: Self = self;
var i: usize = axis;
while (i + 1 < self.rank) : (i += 1) {
view.shape[i] = self.shape[i + 1];
view.strides[i] = self.strides[i + 1];
}
view.shape[self.rank - 1] = 1;
view.strides[self.rank - 1] = 0;
view.rank = @intCast(self.rank - 1);
return view;
}
/// A view with a size-1 axis inserted at `axis`. `axis == rank` appends.
///
/// The new axis gets a stride of 0, which is harmless because its extent is 1 - a
/// single position cannot step anywhere. `isAliased` only objects when extent exceeds 1.
pub fn unsqueeze(self: Self, axis: usize) Error!Self {
if (axis > self.rank or self.rank + 1 > max_rank) {
return Error.OutOfRange;
}
var view: Self = self;
var i: usize = self.rank;
while (i > axis) : (i -= 1) {
view.shape[i] = self.shape[i - 1];
view.strides[i] = self.strides[i - 1];
}
view.shape[axis] = 1;
view.strides[axis] = 0;
view.rank = @intCast(self.rank + 1);
return view;
}
/// A rank-1 view of the whole tensor. Contiguous only, for `reshape`'s reason.
pub fn flatten(self: Self) Error!Self {
return self.reshape(&.{self.size()});
}
/// A view with one axis moved from `from` to `to`, the others keeping their order.
pub fn moveAxis(self: Self, from: usize, to: usize) Error!Self {
if (from >= self.rank or to >= self.rank) {
return Error.OutOfRange;
}
var order: [max_rank]usize = @splat(0);
var k: usize = 0;
var i: usize = 0;
while (i < self.rank) : (i += 1) {
if (i == from) {
continue;
}
order[k] = i;
k += 1;
}
// Open a slot at `to` and drop the moved axis into it.
var j: usize = self.rank - 1;
while (j > to) : (j -= 1) {
order[j] = order[j - 1];
}
order[to] = from;
return self.permute(order[0..self.rank]);
}
/// A view with axes `a` and `b` exchanged. Shares storage; changes through one are visible
/// through the other.
pub fn transpose(self: Self, a: usize, b: usize) Error!Self {
if (a >= self.rank or b >= self.rank) {
return Error.OutOfRange;
}
var view: Self = self;
view.shape[a] = self.shape[b];
view.shape[b] = self.shape[a];
view.strides[a] = self.strides[b];
view.strides[b] = self.strides[a];
return view;
}
};
}
sumAxis/// Sum `a` along `axis` into `out`, which must have `a`'s shape with that axis removed.
///
/// The output's rank is one less, not one with a size-1 axis. Keeping the axis would be the
/// `keepdims` option, and options that change the RESULT SHAPE are better as separate calls than
/// as a flag a reader has to look up - a caller who wants the axis back has `reshape`.
pub fn sumAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void {
comptime requireNumeric(T);
if (axis >= a.rank) {
return Error.OutOfRange;
}
if (out.rank + 1 != a.rank or out.isAliased()) {
return Error.ShapeMismatch;
}
// The surviving axes must match, in order.
var src: usize = 0;
var dst: usize = 0;
while (src < a.rank) : (src += 1) {
if (src == axis) {
continue;
}
if (a.shape[src] != out.shape[dst]) {
return Error.ShapeMismatch;
}
dst += 1;
}
out.fill(0);
var walker: [max_rank]usize = @splat(0);
var reduced: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
var remaining: usize = a.size();
while (remaining > 0) : (remaining -= 1) {
// Project the source position onto the output by dropping the reduced axis.
var k: usize = 0;
var i: usize = 0;
while (i < a.rank) : (i += 1) {
if (i == axis) {
continue;
}
reduced[k] = walker[i];
k += 1;
}
const at: []const usize = reduced[0..out.rank];
const slot: usize = out.offsetOf(at);
out.data[slot] += a.data[a.offsetOf(walker[0..a.rank])];
advance(shape, walker[0..a.rank]);
}
}
The output's rank is one less — the reduced axis is removed, not kept as a size-1
axis. If you want it back, unsqueeze. The surviving extents must match in order.
try zn.meanAxis(f64, out, t, 0);
try zn.maxAxis(f64, out, t, 1);
try zn.minAxis(f64, out, t, 1);
try zn.prodAxis(f64, out, t, 0);
try zn.varianceAxis(f64, out, t, 1, .sample);meanAxis/// Mean along one axis.
pub fn meanAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void {
comptime requireFloat(T);
try sumAxis(T, out, a, axis);
try scale(T, out, out, 1 / @as(T, @floatFromInt(a.shape[axis])));
}
maxAxis/// Largest element along one axis. NaN is not ordered and never wins.
pub fn maxAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void {
comptime requireNumeric(T);
return reduceAxis(T, out, a, axis, zm.lowest(T), struct {
fn f(acc: T, x: T) T {
return if (x > acc) x else acc;
}
}.f);
}
minAxis/// Smallest element along one axis.
pub fn minAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void {
comptime requireNumeric(T);
return reduceAxis(T, out, a, axis, zm.highest(T), struct {
fn f(acc: T, x: T) T {
return if (x < acc) x else acc;
}
}.f);
}
prodAxis/// Product along one axis.
pub fn prodAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void {
comptime requireNumeric(T);
return reduceAxis(T, out, a, axis, 1, struct {
fn f(acc: T, x: T) T {
return acc * x;
}
}.f);
}
varianceAxis/// Variance along one axis, two-pass, with the divisor chosen by `ddof`.
///
/// `out` holds the means after the first pass and is overwritten by the variances in the
/// second, so no scratch the caller has to size. The squared deviations accumulate into the
/// output itself.
pub fn varianceAxis(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
axis: usize,
ddof: Ddof,
) Error!void {
comptime requireFloat(T);
const divisor: f64 = ddof.divisor(a.shape[axis]) orelse return Error.DomainError;
try meanAxis(T, out, a, axis);
// Second pass. Each output slot is visited `shape[axis]` times; on the FIRST visit it still
// holds the mean, which is what the deviation needs - so the mean is read, the slot is
// replaced by the first squared deviation, and later visits add to it. The mean is kept in
// `means` for the remaining visits of that slot.
var means: [max_rank * 512]T = undefined;
if (out.size() > means.len) {
return Error.UnsupportedShape;
}
@memcpy(means[0..out.size()], out.data[out.base..][0..out.size()]);
out.fill(0);
var walker: [max_rank]usize = @splat(0);
var target: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
var remaining: usize = a.size();
while (remaining > 0) : (remaining -= 1) {
var k: usize = 0;
var i: usize = 0;
while (i < a.rank) : (i += 1) {
if (i != axis) {
target[k] = walker[i];
k += 1;
}
}
const slot: usize = out.offsetOf(target[0..out.rank]);
const d: T = (a.data[a.offsetOf(walker[0..a.rank])]) - means[slot - out.base];
out.data[slot] += d * d;
advance(shape, walker[0..a.rank]);
}
try scale(T, out, out, @floatCast(1.0 / divisor));
}
maxAxis, minAxis and prodAxis share one walker with a
comptime accumulation step, so the output-shape check is written once. Each axis reduction is
tested by relating it to its whole-tensor form: reducing along one axis and then over the result
must give maxAll, minAll, prodAll or meanAll,
and varianceAxis on a single row must equal variance of that row.
try zn.cumsum(f64, running, t); // along the last axis
const any = zn.anyNonzero(f64, t);
const all = zn.allNonzero(f64, t); // true for an empty tensor
try zn.where(f64, out, mask, a, b); // a where mask != 0, else b
var order: [n]usize = @splat(0);
try zn.argsort(f64, &order, v); // rank 1, ascending, stablecumsum/// Running sum along the last axis of a rank-1 or rank-2 tensor.
pub fn cumsum(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireNumeric(T);
if (a.rank == 0 or a.rank > 2 or out.rank != a.rank or out.isAliased()) {
return Error.UnsupportedShape;
}
if (out.size() != a.size()) {
return Error.ShapeMismatch;
}
const rows: usize = if (a.rank == 2) a.shape[0] else 1;
const cols: usize = a.shape[a.rank - 1];
var row: usize = 0;
while (row < rows) : (row += 1) {
var running: T = 0;
var col: usize = 0;
while (col < cols) : (col += 1) {
if (a.rank == 2) {
running += try a.at(&.{ row, col });
try out.setAt(&.{ row, col }, running);
} else {
running += try a.at(&.{col});
try out.setAt(&.{col}, running);
}
}
}
}
anyNonzero/// Whether any element is non-zero.
pub fn anyNonzero(comptime T: type, a: Tensor(T)) bool {
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
if (a.data[a.offsetOf(at)] != 0) {
return true;
}
}
return false;
}
allNonzero/// Whether every element is non-zero. True for an empty tensor, by the usual convention.
pub fn allNonzero(comptime T: type, a: Tensor(T)) bool {
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
if (a.data[a.offsetOf(at)] == 0) {
return false;
}
}
return true;
}
where/// `out = a where mask is non-zero, else b`, elementwise. All three broadcast to `out`.
pub fn where(
comptime T: type,
out: Tensor(T),
mask: Tensor(T),
a: Tensor(T),
b: Tensor(T),
) Error!void {
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const vm: Tensor(T) = try mask.broadcastTo(shape);
const va: Tensor(T) = try a.broadcastTo(shape);
const vb: Tensor(T) = try b.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const src: Tensor(T) = if (vm.data[vm.offsetOf(at)] != 0) va else vb;
out.data[out.offsetOf(at)] = src.data[src.offsetOf(at)];
}
}
argsort/// Indices that would sort a rank-1 tensor ascending, written into `out`. Stable.
///
/// Insertion sort: O(n^2), stable, no scratch. A rank-1 sort in a numerics library is almost
/// always over a few hundred elements. `argsort` rather than `sort` because the indices are what
/// a caller usually wants - to reorder something ELSE by this key.
pub fn argsort(comptime T: type, out: []usize, a: Tensor(T)) Error!void {
comptime requireNumeric(T);
if (a.rank != 1) {
return Error.UnsupportedShape;
}
if (out.len != a.shape[0]) {
return Error.ShapeMismatch;
}
for (out, 0..) |*slot, i| {
slot.* = i;
}
var i: usize = 1;
while (i < out.len) : (i += 1) {
const key_index: usize = out[i];
const key: T = try a.at(&.{key_index});
var j: usize = i;
while (j > 0 and (try a.at(&.{out[j - 1]})) > key) : (j -= 1) {
out[j] = out[j - 1];
}
out[j] = key_index;
}
}
where broadcasts all three inputs to out, so a mask from
greater and two full tensors compose directly: selecting by a > b is
maximum(a, b).
argsort returns the permutation rather than the sorted values, because the
indices are what a caller usually wants — to reorder something else by this key. It
is a stable insertion sort: O(n²), no scratch, and a rank-1 sort in a numerics library is
almost always over a few hundred elements. Stability is asserted on a run of equal keys.
try zn.flip(f64, out, a, axis);
try zn.roll(f64, out, a, 1, axis); // positive moves elements FORWARD
const at = try zn.argminAll(f64, a);
const n = try zn.nonzero(f64, &positions, a);
const k = try zn.booleanMask(f64, &values, a, mask);
try zn.sort(f64, &sorted, a);
const distinct = try zn.unique(f64, &out, a);
try zn.bincount(f64, &counts, whole_numbers);flip/// Reverse the order along one axis.
pub fn flip(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void {
comptime requireNumeric(T);
if (axis >= a.rank or out.rank != a.rank or out.isAliased()) {
return Error.UnsupportedShape;
}
for (0..a.rank) |d| {
if (out.shape[d] != a.shape[d]) {
return Error.ShapeMismatch;
}
}
var walker: [max_rank]usize = @splat(0);
var source: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
var remaining: usize = a.size();
while (remaining > 0) : (remaining -= 1) {
@memcpy(source[0..a.rank], walker[0..a.rank]);
source[axis] = a.shape[axis] - 1 - walker[axis];
const value: T = try a.at(source[0..a.rank]);
try out.setAt(walker[0..out.rank], value);
advance(shape, walker[0..a.rank]);
}
}
roll/// Shift elements along one axis, wrapping around: `out[i] = a[(i - shift) mod n]`.
///
/// A POSITIVE SHIFT MOVES ELEMENTS FORWARD, which is numpy's direction and znum's. The sign is
/// the only thing to get wrong here and both conventions look equally reasonable written down,
/// so the test pins it with a shift of 1 on a ramp: the last element must arrive at the front.
///
/// `shift` is signed and may exceed the axis length; the modulus handles both.
pub fn roll(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
shift: isize,
axis: usize,
) Error!void {
comptime requireNumeric(T);
if (axis >= a.rank or out.rank != a.rank or out.isAliased()) {
return Error.UnsupportedShape;
}
for (0..a.rank) |d| {
if (out.shape[d] != a.shape[d]) {
return Error.ShapeMismatch;
}
}
const extent: isize = @intCast(a.shape[axis]);
if (extent == 0) {
return Error.DomainError;
}
var walker: [max_rank]usize = @splat(0);
var source: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
var remaining: usize = a.size();
while (remaining > 0) : (remaining -= 1) {
@memcpy(source[0..a.rank], walker[0..a.rank]);
const here: isize = @intCast(walker[axis]);
source[axis] = @intCast(@mod(here - shift, extent));
const value: T = try a.at(source[0..a.rank]);
try out.setAt(walker[0..out.rank], value);
advance(shape, walker[0..a.rank]);
}
}
argminAll/// The position of the smallest element, scanning flat. Ties go to the first.
///
/// The comparison is `<`, so a NaN never wins: every comparison against NaN is false, and the
/// first non-NaN value keeps the position. The same rule as `argmaxAll`, in the same direction.
pub fn argminAll(comptime T: type, a: Tensor(T)) Error!usize {
comptime requireNumeric(T);
const n: usize = a.size();
if (n == 0) {
return Error.DomainError;
}
var best: usize = 0;
// SEEDED FROM THE FIRST ELEMENT, AND THE NaN CASE IS HANDLED EXPLICITLY
//
// This was `inf(T)`, which made the function float-only. Replacing it with the first element
// looked like a pure simplification and **changed the answer**: seeded from an infinity a NaN
// never won a comparison, so a leading NaN was skipped; seeded from element zero a leading
// NaN IS the seed and nothing can beat it, because every comparison against NaN is false.
//
// The sentinel was not a float artefact - it was doing NaN handling. So the comparison below
// says so directly instead: anything real beats a NaN incumbent. Integers have no NaNs and
// the branch folds away at comptime.
var best_value: T = a.data[0];
var walker: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
for (0..n) |i| {
const x: T = try a.at(walker[0..a.rank]);
if (x < best_value or (comptime isFloat(T)) and zm.isNan(best_value)) {
best_value = x;
best = i;
}
advance(shape, walker[0..a.rank]);
}
return best;
}
nonzero/// The flat positions of every non-zero element, in order, into `out`. Returns how many were
/// written; `Error.ShapeMismatch` if `out` is too short to hold them all.
///
/// The count comes back rather than the caller having to size `out` exactly, because the number
/// of non-zeros is not known until the scan is done and refusing to start would make the caller
/// scan twice.
pub fn nonzero(comptime T: type, out: []usize, a: Tensor(T)) Error!usize {
comptime requireNumeric(T);
var found: usize = 0;
var walker: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
for (0..a.size()) |i| {
if ((try a.at(walker[0..a.rank])) != 0) {
if (found >= out.len) {
return Error.ShapeMismatch;
}
out[found] = i;
found += 1;
}
advance(shape, walker[0..a.rank]);
}
return found;
}
booleanMask/// Copy the elements where `mask` is non-zero into `out`, in order. Returns how many.
pub fn booleanMask(
comptime T: type,
out: []T,
a: Tensor(T),
mask: Tensor(T),
) Error!usize {
comptime requireNumeric(T);
if (mask.size() != a.size() or mask.rank != a.rank) {
return Error.ShapeMismatch;
}
var found: usize = 0;
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
if (mask.data[mask.offsetOf(at)] != 0) {
if (found >= out.len) {
return Error.ShapeMismatch;
}
out[found] = try a.at(at);
found += 1;
}
}
return found;
}
sort/// Sort a copy of `a` ascending into `out`.
pub fn sort(comptime T: type, out: []T, a: Tensor(T)) Error!void {
comptime requireNumeric(T);
if (out.len != a.size()) {
return Error.ShapeMismatch;
}
var walk: Walk = .over(a.shape[0..a.rank]);
var i: usize = 0;
while (walk.next()) |at| : (i += 1) {
out[i] = try a.at(at);
}
std.mem.sort(T, out, {}, std.sort.asc(T));
}
unique/// The distinct values of `a`, ascending, into `out`. Returns how many.
///
/// Sorted rather than first-seen order, because "the distinct values" is a set and a set has no
/// order until one is chosen; ascending is the choice every other library makes and the one that
/// makes two runs comparable.
pub fn unique(comptime T: type, out: []T, a: Tensor(T)) Error!usize {
comptime requireNumeric(T);
if (out.len < a.size()) {
return Error.ShapeMismatch;
}
try sort(T, out[0..a.size()], a);
if (a.size() == 0) {
return 0;
}
var kept: usize = 1;
for (out[1..a.size()]) |value| {
if (value != out[kept - 1]) {
out[kept] = value;
kept += 1;
}
}
return kept;
}
bincount/// Count how many times each whole number appears, into `counts`. Values must be non-negative
/// whole numbers; anything at or past `counts.len` is `Error.OutOfRange`.
///
/// Out of range is an error rather than a silent skip: a bincount that quietly dropped the values
/// it could not hold would return a histogram summing to less than the input, and the caller
/// would have no way to notice.
pub fn bincount(comptime T: type, counts: []usize, a: Tensor(T)) Error!void {
comptime requireNumeric(T);
@memset(counts, 0);
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
const x: T = try a.at(at);
// AN INTEGER IS ALREADY WHOLE, SO THERE IS NOTHING TO CHECK
//
// On a float tensor this has to verify every value is a whole number before using it as
// an index. On an integer tensor the TYPE already said so, and the check would be
// `x != x`. The negative test stays for both - a count of a negative index has no home.
if (comptime isFloat(T)) {
if (x != @floor(x)) {
return Error.DomainError;
}
}
if (x < 0) {
return Error.DomainError;
}
// `@floor` names the rounding for a float; an integer is already the index.
const slot: usize = if (comptime isFloat(T)) @floor(x) else @intCast(x);
if (slot >= counts.len) {
return Error.OutOfRange;
}
counts[slot] += 1;
}
}
Every function here has a convention that could reasonably go the other way, so each is pinned rather than left for a reader to run it. A positive roll moves elements
forward, so roll(1) brings the last value to the front — numpy's
direction. unique returns ascending order, not first-seen; on
the sample 3, 1, 4, 1, 5, 3, 1 the two conventions differ.
nonzero and booleanMask are the same question asked two ways
— the positions, and the values at those positions — and they agree
with each other rather than against two lists written by hand.
Both return a count rather than requiring the caller to size the output exactly, because the
number of non-zeros is not known until the scan is done and refusing to start would make the
caller scan twice. An output too small to hold them all is ShapeMismatch, never a
truncation. bincount takes the same line: a value past the end of
counts is OutOfRange, because a count that quietly dropped what it could
not hold would sum to less than the input with no way to notice.
const total = try zn.logSumExp(f64, a);
try zn.logSumExpAxis(f64, per_row, a, 1); // one maximum per row
try zn.stdDevAxis(f64, out, a, 1, .sample);logSumExp/// `log(sum(exp(a)))`, computed so it does not overflow.
///
/// SUBTRACTING THE MAXIMUM IS THE WHOLE FUNCTION
///
/// The direct reading overflows almost immediately: `exp(800)` is infinity in f64, so a sum of
/// exponentials of ordinary-looking numbers comes back as `inf` and the logarithm of that is
/// `inf` too. Subtracting the largest element first makes the biggest term exactly `exp(0) = 1`
/// and every other term smaller, so nothing can overflow - and adding the maximum back at the
/// end is exact, because `log(sum(exp(a - m))) + m` IS `log(sum(exp(a)))` by algebra.
///
/// a = { 800, 801, 802 } direct: inf this: 802.4076
///
/// The same trick is what makes `softmaxRows` safe, and this is the reduction underneath cross
/// entropy: `crossEntropy = logSumExp(logits) - logits[correct]`.
///
/// An empty tensor has no maximum to subtract and the sum of no terms is zero, whose logarithm
/// is negative infinity - which is the right answer but not one a caller usually wants by
/// accident, so it is a `DomainError`.
pub fn logSumExp(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
if (a.size() == 0) {
return Error.DomainError;
}
const largest: T = try maxAll(T, a);
// An infinite maximum has no finite offset to remove; the answer is that infinity.
if (!zm.isFinite(largest)) {
return largest;
}
var total: CompensatedSum(T) = .{};
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
total.add(@exp((try a.at(at)) - largest));
}
return largest + @log(total.value());
}
logSumExpAxis/// `logSumExp` along one axis of a rank-2 tensor, written into `out`.
///
/// Each row (or column) gets its own maximum, which is what makes this useful for a batch: one
/// row saturating does not affect the stability of its neighbours.
pub fn logSumExpAxis(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
axis: usize,
) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or out.rank != 1 or axis > 1) {
return Error.UnsupportedShape;
}
const kept: usize = a.shape[1 - axis];
const along: usize = a.shape[axis];
if (out.shape[0] != kept or along == 0) {
return Error.ShapeMismatch;
}
for (0..kept) |i| {
var largest: T = -inf(T);
for (0..along) |j| {
const at: []const usize = if (axis == 0) &.{ j, i } else &.{ i, j };
largest = @max(largest, try a.at(at));
}
if (!zm.isFinite(largest)) {
try out.setAt(&.{i}, largest);
continue;
}
var total: CompensatedSum(T) = .{};
for (0..along) |j| {
const at: []const usize = if (axis == 0) &.{ j, i } else &.{ i, j };
total.add(@exp((try a.at(at)) - largest));
}
try out.setAt(&.{i}, largest + @log(total.value()));
}
}
stdDevAxis/// Standard deviation along one axis: the square root of `varianceAxis`.
///
/// Thin on purpose. It exists because `@sqrt(varianceAxis(...))` needs a second buffer and a
/// second pass at every call site, and because a reader looking for a standard deviation should
/// find one rather than have to know it is a square root away.
pub fn stdDevAxis(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
axis: usize,
ddof: Ddof,
) Error!void {
comptime requireFloat(T);
try varianceAxis(T, out, a, axis, ddof);
for (out.data) |*slot| {
slot.* = @sqrt(slot.*);
}
}
The direct reading of log(sum(exp(a))) overflows almost immediately —
exp(800) is infinity in f64, so a sum of exponentials of ordinary-looking numbers
comes back inf and its logarithm is inf too:
a = { 800, 801, 802 } direct: inf logSumExp: 802.4076Subtracting the largest element first is the whole function. It makes the
biggest term exactly exp(0) = 1 and every other one smaller, so nothing can overflow —
and adding the maximum back at the end is exact, because log(sum(exp(a - m))) + m
is log(sum(exp(a))) by algebra, not by approximation.
The test asserts three things rather than one: that the direct form really does produce
inf (so the function has a reason to exist), that this one agrees with the direct form
to 1e-14 where the direct form works (so the trick costs nothing), and that shifting every
element by 500 shifts the answer by exactly 500 (the identity the subtraction relies on — break
the subtraction and this breaks).
logSumExpAxis takes a maximum per row, so one saturating row does not destabilise
its neighbours. There is a GPU twin, log_sum_exp_rows, one thread per
row — and it subtracts the maximum for the same reason softmax_rows beside it does.
A shader has no more headroom than a host: exp(800) is infinity in f32 long before f64
gives up. One trick, two kernels, and if either drops it the failure is inf rather
than a wrong digit. This is the reduction underneath cross entropy:
crossEntropy = logSumExp(logits) - logits[correct].
const d = try zn.determinant(f64, a, &pivots); // OVERWRITES `a`determinant/// The determinant, by LU factorisation.
///
/// THE PIVOT SIGN IS WHY THIS IS FIVE LINES AND NOT FIFTY
///
/// `lu` already does the hard part and already returns the sign the row swaps accumulated, so
/// the determinant is that sign times the product of the diagonal. Writing a separate
/// cofactor expansion would be a second implementation of the same arithmetic, `O(n!)` instead
/// of `O(n^3)`, and a second place for a bug.
///
/// **`a` is overwritten** - `lu` factorises in place, which is what makes it cheap. A caller who
/// still needs the matrix passes a copy, and the doc says so rather than the caller discovering
/// it.
///
/// A SINGULAR MATRIX IS ZERO HERE AND AN ERROR IN `lu`, AND BOTH ARE RIGHT
///
/// `lu` returns `DomainError` when it cannot pivot, because a factorisation of a singular matrix
/// is not usable for what factorisations are for - solving. A determinant is different: **the
/// determinant of a singular matrix is exactly zero**, which is a complete and useful answer, and
/// is precisely how a caller tests for singularity in the first place.
///
/// So the error is caught here and turned into the zero it means. The same underlying condition
/// is a failure for one caller and an answer for another, and the function each of them calls
/// should say which.
pub fn determinant(comptime T: type, a: Tensor(T), pivots: []usize) Error!T {
comptime requireFloat(T);
const pivot_sign: T = lu(T, a, pivots) catch |failure| switch (failure) {
Error.DomainError => return 0,
else => return failure,
};
var product: T = pivot_sign;
for (0..a.shape[0]) |i| {
product *= try a.at(&.{ i, i });
}
return product;
}
lu already does the hard part and already returns the sign the row swaps
accumulated, so the determinant is that sign times the diagonal — five lines rather than a
cofactor expansion, which would be a second implementation of the same arithmetic at
O(n!) instead of O(n³).
The test swaps two rows and checks the answer negates. A version that lost the pivot sign would return −306 both times and pass a single-matrix test looking perfectly correct.
A singular matrix is zero here and an error in lu, and both are right.
lu returns DomainError when it cannot pivot, because a factorisation of a
singular matrix is not usable for what factorisations are for. A determinant is different: the
determinant of a singular matrix is exactly zero, which is a complete answer and is precisely how
a caller tests for singularity. The same underlying condition is a failure for one caller and an
answer for another, and each function says which.
zn.all(i32, t) every element nonzero
zn.any(i32, t) some element nonzero
zn.countNonzero(i32, t) how manyall/// True when every element is nonzero. An empty tensor is TRUE.
///
/// THE EMPTY CASE IS NOT ARBITRARY
///
/// `all` over nothing is true and `any` over nothing is false, and that is not a coin flip - it
/// is what makes them compose. `all(a ++ b)` must equal `all(a) and all(b)` for every pair
/// including an empty one, which forces `all(empty) == true`; the same argument on `or` forces
/// `any(empty) == false`.
///
/// numpy, Python's builtins and SQL all agree, so a reader can check an answer against any of
/// them. Getting it backwards produces a function that looks right on every non-empty test.
pub fn all(comptime T: type, a: Tensor(T)) bool {
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
if (a.data[a.offsetOf(at)] == 0) {
return false;
}
}
return true;
}
any/// True when any element is nonzero. An empty tensor is FALSE - see `all`.
pub fn any(comptime T: type, a: Tensor(T)) bool {
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
if (a.data[a.offsetOf(at)] != 0) {
return true;
}
}
return false;
}
countNonzero/// How many elements are nonzero.
///
/// Returns a `usize`, not a `T`. A count is not a measurement in the tensor's units, and at f32
/// it would stop being exact at 2^24 - the same ceiling the index-tensor work ran into.
pub fn countNonzero(comptime T: type, a: Tensor(T)) usize {
var total: usize = 0;
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
if (a.data[a.offsetOf(at)] != 0) {
total += 1;
}
}
return total;
}
all over nothing is true. any over nothing is
false. That looks arbitrary and is not:
all(a ++ b) == all(a) and all(b) must hold for every pair
take b empty -> all(a) == all(a) and all(empty)
so all(empty) must be trueThe same argument on or forces any(empty) == false. numpy, Python's
builtins and SQL all agree, so an answer can be checked against any of them.
Getting it backwards produces a function that looks right on every non-empty test, the composition directly rather than just the two values.
values 1, NaN, 3
sumAll -> NaN correct: it tells you one arrived
nanSum -> 4 skippedsumAll propagating a NaN is useful. Skipping silently by default would
hide it, so the skipping version is a different name and reaching for it is a decision the caller
makes rather than a default they inherit.
For a column with real gaps, Series.sumValid is better still: it skips what is
absent rather than what is NaN — and for an integer column there is no NaN to
skip.
One small consistency: nanMean of an all-NaN tensor returns NaN, where
Series.meanValid returns null. The difference is deliberate — here the caller
already has NaNs in their data; there the absence was recorded rather than inferred from a
value.
zn.argExtreme(f64, .largest, t) which element wins, ANYWHERE -> one number
zn.argmaxAxis(f64, i32, out, t, 1) which column wins in each row -> a tensorargExtreme/// The FLAT index of the largest or smallest element in the whole tensor.
///
/// FLAT, WHERE `argmaxAxis` IS PER-AXIS - AND THAT IS THE WHOLE DIFFERENCE
///
/// `argmaxAxis` answers "which column wins in each row". This answers "which element wins,
/// anywhere". Both are useful and they are not variants of each other: one returns a tensor of
/// indices, the other a single number.
///
/// The number is a FLAT index into the walk order, not a coordinate. A coordinate would need a
/// buffer to return it in, and the caller who wants one has `unravel` arithmetic that depends on
/// a shape they already have.
///
/// TIES GO TO THE FIRST, which is numpy's rule and the only one that makes the answer independent
/// of how the tensor is traversed. A later-wins rule would give a different index for a tensor and
/// its own transpose.
///
/// `which` is comptime, so the comparison is one branch resolved at compile time rather than a
/// test in the inner loop.
pub fn argExtreme(comptime T: type, comptime which: Extreme, a: Tensor(T)) Error!usize {
comptime requireNumeric(T);
if (a.size() == 0) {
return Error.DomainError;
}
var best: T = if (which == .largest) zm.lowest(T) else zm.highest(T);
var best_at: usize = 0;
var position: usize = 0;
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| : (position += 1) {
const x: T = a.data[a.offsetOf(at)];
const better: bool = if (which == .largest) x > best else x < best;
if (better) {
best = x;
best_at = position;
}
}
return best_at;
}
argmaxAxis/// The index of the largest element along `axis`, written into `out` as a float.
///
/// A float, because every tensor here holds one type and an index tensor would need a second.
/// Every index a rank-2 tensor can address is exact in f32 up to 2^24, and in f64 up to 2^53 -
/// far past any axis this library can hold. Stated rather than left to be discovered.
pub fn argmaxAxis(
comptime T: type,
comptime Index: type,
out: Tensor(Index),
a: Tensor(T),
axis: usize,
) Error!void {
comptime requireNumeric(T);
comptime requireNumeric(Index);
return argExtremumAxis(T, Index, out, a, axis, true);
}
These are not variants of each other. One returns a single index, the other a tensor of them, and reaching for the wrong one is a shape error rather than a wrong answer — which is at least loud.
The index is flat, a position in the walk, not an offset into the buffer. That distinction is what makes the answer usable on a strided column: the middle column of a matrix has its largest element at walk position 0 while sitting at buffer index 1.
Ties go to the first. numpy's rule, and the only one that makes the answer independent of traversal order — a later-wins rule would give a different index for a tensor and its own transpose, which is the kind of inconsistency nobody thinks to look for.
Most reduction rows in the GPU sweep carry a ULP allowance, because a sum accumulates differently depending on the order the threads finish in. Two of the new ones do not:
count nonzero bar 0 a count INCREMENTS; integer addition is exact in any order
diff bar 0 one subtraction, one rounding, on both sidesThat is worth knowing when reading the sweep: a nonzero difference in those rows is a real bug, not rounding. Choosing a generous bar for every reduction because most of them need one would have thrown that away.
Every operation in this library has a GPU kernel checked against its CPU twin on the same input. Each row carries a tolerance, and choosing it is the interesting part.
all nonzero tol 0 exact
any nonzero tol 0 exact
count nonzero tol 0 exact
diff tol 0 exact
prod all tol 1e-3 NOT exact
sum all tol 6e-4 NOT exactall and any combine with and and or, which
are associative and idempotent. No order of combination can change the answer, so
any difference at all is a real bug rather than accumulated rounding.
count increments an integer. diff subtracts neighbouring pairs, each
rounding once and independently. Same story: the arithmetic has no room to drift.
A zero bar is a stronger claim than a small one, and using one where it holds is how the sweep stays informative. A tolerance everywhere would pass a kernel that was subtly wrong.
Multiplication is associative in real arithmetic and is not in floating
point: (a*b)*c and a*(b*c) can differ in the last bits. A product
accumulates that over the whole buffer, where a sum accumulates only absolute error — and it
reaches infinity far sooner.
So the tolerance beside two zeros is not laziness about the product. It is the measurement, and the contrast is what makes all three bars mean something.
all_nonzero exits early on the first zero. A kernel that exited on the first
nonzero instead would pass any test that only fed it an all-zero buffer. The sweep runs
the real field — a spread of values in an order nobody chose — against the host
function, which is how an early-exit mistake surfaces.
const area = try zn.trapz(f64, y, spacing); evenly spaced
const area = try zn.trapzCoords(f64, y, x); positions giventrapz/// The integral of a sampled function by the trapezoid rule, evenly spaced.
///
/// WHY TRAPEZOIDS AND NOT RECTANGLES
///
/// A rectangle takes one endpoint of each interval and ignores the other, so it is systematically
/// wrong on any sloped function - too low on a rise, too high on a fall. A trapezoid averages the
/// two, which makes it **exact for any straight line** and therefore exact for the two cases
/// people check first.
///
/// The weights fall out of that: the interior samples each belong to two intervals and count
/// once, while the two ENDPOINTS belong to one each and count half. Getting that backwards gives
/// an answer wrong by roughly one sample's worth, which on a long series looks like a small
/// calibration error rather than a bug.
pub fn trapz(comptime T: type, y: Tensor(T), spacing: T) Error!T {
comptime requireFloat(T);
const n: usize = y.size();
if (n == 0) {
return Error.DomainError;
}
if (n == 1) {
// One sample spans no interval, so the integral is zero rather than the sample's value.
return 0;
}
var total: CompensatedSum(T) = .{};
for (0..n) |i| {
const weight: T = if (i == 0 or i == n - 1) 0.5 else 1.0;
total.add(weight * y.data[y.offsetOf(&.{i})]);
}
return total.value() * spacing;
}
trapzCoords/// The same, with the sample positions given rather than assumed even.
///
/// Each interval carries its own width, so this handles data logged whenever something happened
/// rather than on a clock. `x` must be increasing: a repeated or decreasing coordinate would
/// contribute a zero or negative width, and silently subtracting area is worse than refusing.
pub fn trapzCoords(comptime T: type, y: Tensor(T), x: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = y.size();
if (n != x.size()) {
return Error.ShapeMismatch;
}
if (n == 0) {
return Error.DomainError;
}
if (n == 1) {
return 0;
}
var total: CompensatedSum(T) = .{};
for (0..n - 1) |i| {
const left: T = x.data[x.offsetOf(&.{i})];
const right: T = x.data[x.offsetOf(&.{i + 1})];
if (!(right > left)) {
return Error.DomainError;
}
const ly: T = y.data[y.offsetOf(&.{i})];
const ry: T = y.data[y.offsetOf(&.{i + 1})];
total.add(0.5 * (ly + ry) * (right - left));
}
return total.value();
}
A rectangle rule takes one endpoint of each interval and ignores the other, so it is systematically wrong on any slope — too low on a rise, too high on a fall. A trapezoid averages the two, which makes it exact for any straight line.
y = x over [0,4], spacing 1 true 8
trapezoid 8 exact
left rect 6 too low
right rect 10 too highThis is the classic error. Five samples at unit spacing span four intervals:
y = 1, five samples -> integral 4, not 5The interior samples each belong to two intervals and count once; the two endpoints belong to one each and count half. Weighting everything equally gives 5 — wrong by exactly one sample's worth, which on a long series reads like a small calibration error rather than a bug.
A consequence worth knowing: one sample integrates to zero, not to its own value. It spans no interval.
On a curve the answer is approximate and the error has a known sign: trapezoids over-estimate a convex function, because the chord sits above the curve. The test asserts that direction, not just the magnitude.
trapzCoords handles data logged whenever something happened rather than on a clock.
A repeated coordinate contributes zero width and a decreasing one would
subtract area — silently. Both are refused.
And the two agree on the case they share: even coordinates through trapzCoords give
exactly what trapz gives and so the implementations cannot
drift.
rfft returns n/2 + 1 bins and says nothing about what frequency each
one is. The mapping is k / (n * spacing):
8 samples, 1 ms apart -> 0, 125, 250, 375, 500 HzBin 0 is DC and the last is Nyquist — the highest frequency a sample rate
can represent. Working it out by hand goes right until n is odd, and the symptom is a
plot whose peak sits at the wrong pitch.
try zn.matmul(f64, out, a, b); // (m,k) @ (k,n) -> (m,n)
try zn.matmul(f64, out, a, try b.transpose(0, 1)); // a @ bT, no copymatmul/// `out = a @ b` - the matrix product, for rank-2 tensors only.
///
/// THE ONE OPERATION THAT IS NOT ELEMENTWISE
///
/// `(m, k) @ (k, n) -> (m, n)`. Every shape is checked; a mismatch is `ShapeMismatch` before any
/// element is written, so a partially-filled output is not a state this can produce.
///
/// The inner sum is compensated, for the reason in `sumAll`: a dot product over a long
/// contraction axis is exactly the accumulation that loses its low bits, and in a neural network
/// that axis is the layer width. Getting this wrong does not crash - it makes gradients slightly
/// wrong, everywhere, in a way that reads as "training is a bit unstable".
///
/// This is the naive triple loop, deliberately. It is the SPECIFICATION every faster
/// implementation is checked against, and it stays in the file when a blocked or tiled version
/// arrives. Strided and stretched inputs are handled by going through `at`, so a transposed view
/// multiplies correctly without being copied first.
pub fn matmul(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or b.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
const m: usize = a.shape[0];
const k: usize = a.shape[1];
const n: usize = b.shape[1];
if (b.shape[0] != k or out.shape[0] != m or out.shape[1] != n) {
return Error.ShapeMismatch;
}
if (out.isAliased()) {
return Error.UnsupportedShape;
}
var row: usize = 0;
while (row < m) : (row += 1) {
var col: usize = 0;
while (col < n) : (col += 1) {
var total: CompensatedSum(T) = .{};
var i: usize = 0;
while (i < k) : (i += 1) {
const product: T = a.at2(row, i) *
b.at2(i, col);
total.add(product);
}
out.setAt2(row, col, total.value());
}
}
}
Rank 2 only. Every shape is checked before any element is written, so a partially-filled output is not a state this can produce. The inner sum is compensated.
Strided inputs work without being copied, because the implementation reads through
at. A transposed operand is therefore free on the CPU — a stride swap —
which is not true of the GPU kernels, where the transpose has to live in the index arithmetic or
the operand has to be materialised first.
This is the naive triple loop. It is the specification a faster implementation is checked against.
const d = try zn.dotAll(f64, u, v); // sizes must match; shapes need not
const tr = try zn.trace(f64, a); // square, rank 2
const nm = try zn.norm(f64, a); // Frobenius, over every elementdotAll/// The dot product of two tensors with the same number of elements.
///
/// Shapes need not match, only sizes: a row and a column of the same length have the same dot
/// product, and requiring identical shapes would force a reshape at every call site.
pub fn dotAll(comptime T: type, a: Tensor(T), b: Tensor(T)) Error!T {
comptime requireNumeric(T);
const n: usize = a.size();
if (n != b.size()) {
return Error.ShapeMismatch;
}
var total: CompensatedSum(T) = .{};
var pair: Walk2 = .over(a.shape[0..a.rank], b.shape[0..b.rank]);
while (pair.next()) |at| {
const term: T = (try a.at(at.first)) *
(try b.at(at.second));
total.add(term);
}
return total.value();
}
trace/// The sum of the diagonal of a square rank-2 tensor.
pub fn trace(comptime T: type, a: Tensor(T)) Error!T {
comptime requireNumeric(T);
if (a.rank != 2 or a.shape[0] != a.shape[1]) {
return Error.UnsupportedShape;
}
var total: T = 0;
var i: usize = 0;
while (i < a.shape[0]) : (i += 1) {
total += a.at2(i, i);
}
return total;
}
norm/// The Euclidean (Frobenius) norm: the square root of the sum of squares.
///
/// SCALED BY THE LARGEST MAGNITUDE BEFORE SQUARING. A vector whose elements are near
/// `sqrt(maxFloat)` overflows a direct sum of squares even though its norm is representable, and
/// one near `sqrt(minFloat)` underflows to zero. Factoring the largest element out costs one extra
/// pass and removes both failures.
pub fn norm(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
var largest: T = 0;
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
largest = @max(largest, @abs(try a.at(at)));
}
if (largest == 0) {
return 0;
}
// A SECOND pass over the same shape: a fresh `Walk` rather than resetting the first, which
// is what the old code did and is one more thing to forget.
var second: Walk = .over(a.shape[0..a.rank]);
var total: T = 0;
while (second.next()) |at| {
const scaled: T = (try a.at(at)) / largest;
total += scaled * scaled;
}
return largest * @sqrt(total);
}
dotAll requires equal element counts rather than equal shapes, so a row and a column
of the same length pair without a reshape. Its accumulation is compensated.
norm divides by the largest magnitude before squaring and multiplies back
afterwards:
if (largest == 0) return 0;
// ... total += (x / largest) * (x / largest) ...
return largest * @sqrt(total);A direct sum of squares overflows for a vector whose elements approach √maxFloat and underflows to zero near √minFloat, in both cases for inputs whose norm is perfectly representable. The extra pass removes both. A vector of two 1e200 values returns 1.414e200 rather than infinity.
var pivots: [3]usize = @splat(0);
const sign = try zn.lu(f64, a, &pivots); // a is overwritten
try zn.solve(f64, a, b, &pivots); // a overwritten, b becomes x
const det = try zn.determinantLu(f64, a, &pivots); // a overwrittenlu/// An LU factorisation with partial pivoting, computed in place.
///
/// WHAT `lu` LEAVES BEHIND
///
/// `a` is overwritten: the strict lower triangle holds L's off-diagonal entries (L's diagonal is
/// all ones and is not stored), and the upper triangle including the diagonal holds U.
/// `pivots[i]` is the row swapped into position `i`. The return value is the SIGN of the
/// permutation, +1 or -1, which `determinant` needs and which would otherwise have to be
/// recomputed by counting the swaps.
///
/// PARTIAL PIVOTING IS NOT OPTIONAL. Without it a matrix whose leading entry is small but
/// non-zero divides by that small number and loses most of its precision, and a leading zero
/// fails outright on a matrix that is perfectly well conditioned. Choosing the largest available
/// pivot in the column bounds the multipliers by 1.
///
/// A column with no non-zero pivot means the matrix is singular: `DomainError`, rather than an
/// infinity that reaches the caller as a plausible-looking result.
pub fn lu(comptime T: type, a: Tensor(T), pivots: []usize) Error!T {
comptime requireFloat(T);
if (a.rank != 2 or a.shape[0] != a.shape[1]) {
return Error.UnsupportedShape;
}
const n: usize = a.shape[0];
if (pivots.len != n) {
return Error.ShapeMismatch;
}
var pivot_sign: T = 1;
var col: usize = 0;
while (col < n) : (col += 1) {
// Largest magnitude at or below the diagonal becomes the pivot.
var best: usize = col;
var best_value: T = @abs(a.at2(col, col));
var row: usize = col + 1;
while (row < n) : (row += 1) {
const candidate: T = @abs(a.at2(row, col));
if (candidate > best_value) {
best_value = candidate;
best = row;
}
}
if (best_value == 0) {
return Error.DomainError;
}
pivots[col] = best;
if (best != col) {
pivot_sign = -pivot_sign;
var j: usize = 0;
while (j < n) : (j += 1) {
const top: T = a.at2(col, j);
a.setAt2(col, j, a.at2(best, j));
a.setAt2(best, j, top);
}
}
const pivot: T = a.at2(col, col);
row = col + 1;
while (row < n) : (row += 1) {
const factor: T = a.at2(row, col) / pivot;
a.setAt2(row, col, factor);
var j: usize = col + 1;
while (j < n) : (j += 1) {
const updated: T = a.at2(row, j) -
factor * a.at2(col, j);
a.setAt2(row, j, updated);
}
}
}
return pivot_sign;
}
solve/// Solve `A x = b` for a square `A`, overwriting `a` with its factorisation and `b` with `x`.
///
/// `b` may have several columns, each an independent right-hand side. Solving for many at once
/// reuses the one factorisation, which is the expensive part.
pub fn solve(comptime T: type, a: Tensor(T), b: Tensor(T), pivots: []usize) Error!void {
comptime requireFloat(T);
if (b.rank != 2 or b.shape[0] != a.shape[0]) {
return Error.ShapeMismatch;
}
_ = try lu(T, a, pivots);
const n: usize = a.shape[0];
const rhs: usize = b.shape[1];
// Apply the same row swaps to b.
var i: usize = 0;
while (i < n) : (i += 1) {
if (pivots[i] != i) {
var j: usize = 0;
while (j < rhs) : (j += 1) {
const top: T = b.at2(i, j);
b.setAt2(i, j, b.at2(pivots[i], j));
b.setAt2(pivots[i], j, top);
}
}
}
// Forward substitution through L, whose diagonal is implicitly 1.
i = 0;
while (i < n) : (i += 1) {
var j: usize = 0;
while (j < rhs) : (j += 1) {
var acc: T = b.at2(i, j);
var k: usize = 0;
while (k < i) : (k += 1) {
acc -= a.at2(i, k) * b.at2(k, j);
}
b.setAt2(i, j, acc);
}
}
// Back substitution through U.
i = n;
while (i > 0) {
i -= 1;
var j: usize = 0;
while (j < rhs) : (j += 1) {
var acc: T = b.at2(i, j);
var k: usize = i + 1;
while (k < n) : (k += 1) {
acc -= a.at2(i, k) * b.at2(k, j);
}
b.setAt2(i, j, acc / a.at2(i, i));
}
}
}
determinantLu/// The determinant of a square matrix. `a` is overwritten by its factorisation.
///
/// The product of U's diagonal times the permutation sign. Returns 0 for a singular matrix
/// rather than propagating `lu`'s `DomainError`: a determinant of zero is the correct answer to
/// the question, not a failure to answer it.
pub fn determinantLu(comptime T: type, a: Tensor(T), pivots: []usize) Error!T {
comptime requireFloat(T);
const permutation_sign: T = lu(T, a, pivots) catch |err| switch (err) {
Error.DomainError => return 0,
else => return err,
};
var product: T = permutation_sign;
var i: usize = 0;
while (i < a.shape[0]) : (i += 1) {
product *= a.at2(i, i);
}
return product;
}
lu factorises in place with partial pivoting. Afterwards the strict lower triangle
of a holds L's off-diagonal entries — L's diagonal is all ones and is not stored —
and the upper triangle including the diagonal holds U. pivots[i] is the row swapped
into position i. The return value is the sign of the permutation, +1 or −1,
which determinantLu needs.
Pivoting selects the largest available magnitude in the column, which bounds every multiplier by 1. Without it, a small leading entry divides by that small number and loses most of the available precision, and a zero leading entry fails on a matrix that is otherwise well conditioned.
A column with no non-zero pivot means the matrix is singular. lu and
solve return DomainError; determinantLu returns
0, because a determinant of zero is the answer to the question rather than a
failure to answer it.
solve factorises, applies the same row swaps to b, then substitutes
forward through L and back through U. b may have several columns, each an independent
right-hand side; solving for many at once reuses the single factorisation, which is the expensive
part.
All three overwrite a. Keep a copy if you need the original — checking a solution's
residual requires one.
try zn.qr(f64, q, r, a); // a is (m, n) with m >= n; q is (m, m), r is (m, n)qr/// QR by Householder reflections: `a = q * r`, with `q` orthogonal and `r` upper triangular.
///
/// `a` is `(m, n)` with `m >= n`; `q` is `(m, m)` and `r` is `(m, n)`. Both are overwritten.
///
/// HOUSEHOLDER, NOT GRAM-SCHMIDT, AND THE DIFFERENCE IS MEASURED
///
/// Gram-Schmidt builds `q` a column at a time by subtracting projections. Even the modified form
/// loses orthogonality in proportion to the matrix's condition number: on an ill-conditioned
/// input the columns it returns are **visibly not orthogonal**, and every later use of `q`
/// inherits that error. Householder builds `q` from reflections, each orthogonal to working
/// precision by construction, so their product is too however badly conditioned the input.
///
/// The test runs both on the same matrix and asserts the gap - this is not a preference
/// between algorithms, it is one of them being wrong.
///
/// Each reflection zeroes one column below the diagonal. `alpha` takes the sign OPPOSITE the
/// pivot so `v = x - alpha*e1` never subtracts two nearby numbers; the same sign would cancel
/// catastrophically exactly when the pivot dominates its column, which is the common case.
pub fn qr(comptime T: type, q: Tensor(T), r: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or q.rank != 2 or r.rank != 2) {
return Error.UnsupportedShape;
}
const m: usize = a.shape[0];
const n: usize = a.shape[1];
if (m < n) {
return Error.UnsupportedShape;
}
if (q.shape[0] != m or q.shape[1] != m or r.shape[0] != m or r.shape[1] != n) {
return Error.ShapeMismatch;
}
var reflector: [max_rank * 512]T = undefined;
if (m > reflector.len) {
return Error.UnsupportedShape;
}
try materialise(T, r, a);
q.fill(0);
var d: usize = 0;
while (d < m) : (d += 1) {
try q.setAt(&.{ d, d }, 1);
}
var k: usize = 0;
while (k < n) : (k += 1) {
var norm_squared: CompensatedSum(T) = .{};
var i: usize = k;
while (i < m) : (i += 1) {
const x: T = try r.at(&.{ i, k });
norm_squared.add(x * x);
}
const column_norm: T = @sqrt(norm_squared.value());
if (column_norm == 0) {
continue;
}
const pivot: T = try r.at(&.{ k, k });
const alpha: T = if (pivot > 0) -column_norm else column_norm;
reflector[k] = pivot - alpha;
i = k + 1;
while (i < m) : (i += 1) {
reflector[i] = try r.at(&.{ i, k });
}
var v_squared: CompensatedSum(T) = .{};
i = k;
while (i < m) : (i += 1) {
v_squared.add(reflector[i] * reflector[i]);
}
const v_norm_squared: T = v_squared.value();
if (v_norm_squared == 0) {
continue;
}
var j: usize = k;
while (j < n) : (j += 1) {
var projection: CompensatedSum(T) = .{};
i = k;
while (i < m) : (i += 1) {
projection.add(reflector[i] * (try r.at(&.{ i, j })));
}
const factor: T = 2 * projection.value() / v_norm_squared;
i = k;
while (i < m) : (i += 1) {
try r.setAt(&.{ i, j }, (try r.at(&.{ i, j })) - factor * reflector[i]);
}
}
var row: usize = 0;
while (row < m) : (row += 1) {
var projection: CompensatedSum(T) = .{};
i = k;
while (i < m) : (i += 1) {
projection.add(reflector[i] * (try q.at(&.{ row, i })));
}
const factor: T = 2 * projection.value() / v_norm_squared;
i = k;
while (i < m) : (i += 1) {
try q.setAt(&.{ row, i }, (try q.at(&.{ row, i })) - factor * reflector[i]);
}
}
}
// The reflections leave arithmetic dust below the diagonal; it is structurally zero.
var row: usize = 1;
while (row < m) : (row += 1) {
var col: usize = 0;
while (col < @min(row, n)) : (col += 1) {
try r.setAt(&.{ row, col }, 0);
}
}
}
Gram–Schmidt builds q a column at a time by subtracting projections, and even
the modified form loses orthogonality in proportion to the matrix's condition number. Householder
builds q from reflections, each orthogonal to working precision by construction, so
their product is too however badly conditioned the input. Measured on Hilbert matrices —
ill-conditioned by construction, one line to build:
| Size | Householder | Gram–Schmidt |
|---|---|---|
| 6×6 | 4.4e-16 | 1.8e-10 |
| 8×8 | 3.3e-16 | 4.4e-7 |
| 10×10 | 4.4e-16 | 1.8e-4 |
Householder does not degrade with size at all. A 1.8e-4 error is not a rounding difference
— it is q not being orthogonal, and every later use of it inherits that.
The test asserts Gram–Schmidt is visibly worse, so if that ever stopped being true
the reason for the extra work would be gone.
Each reflection takes the sign opposite the pivot when forming
v = x − αe₁, so it never subtracts two nearby numbers; the same sign
cancels catastrophically exactly when the pivot dominates its column, which is the common case.
r is upper triangular exactly — the dust below the diagonal is zeroed,
not merely small.
try zn.eigh(f64, &values, vectors, a); // a is overwritten; vectors' COLUMNS are eigenvectorseigh/// Eigenvalues and eigenvectors of a symmetric matrix, by cyclic Jacobi rotations.
///
/// `a` is `(n, n)` and symmetric, and is OVERWRITTEN - it becomes the diagonal matrix. `values`
/// receives the eigenvalues, `vectors` the eigenvectors as its COLUMNS, so `a*vectors[:,k] =
/// values[k]*vectors[:,k]`.
///
/// THE CONVERGENCE TEST IS RELATIVE TO THE MATRIX, NOT AN ABSOLUTE CONSTANT
///
/// Each rotation zeroes one off-diagonal pair and the sweep repeats until what remains off the
/// diagonal is negligible. **Negligible compared to what** is the whole question. A fixed
/// threshold - znum uses `1e-30` on the sum of squares - declares any sufficiently small matrix
/// already diagonal: scale a matrix by `1e-16` and its off-diagonal sum of squares falls below
/// `1e-30` before a single rotation, so the routine returns the untouched diagonal as the
/// eigenvalues. They are not the eigenvalues, and nothing reports a problem.
///
/// Measuring against the matrix's own Frobenius norm makes the test scale-free: the same
/// matrix times `1e-20` gives the same eigenvalues times `1e-20`, which the test asserts.
///
/// `DomainError` if the sweeps run out. Cyclic Jacobi converges quadratically and reaches
/// working precision in well under ten sweeps for any matrix this will see; returning an
/// unconverged answer silently is the failure this avoids.
pub fn eigh(comptime T: type, values: []T, vectors: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or vectors.rank != 2) {
return Error.UnsupportedShape;
}
const n: usize = a.shape[0];
if (a.shape[1] != n or vectors.shape[0] != n or vectors.shape[1] != n or values.len != n) {
return Error.ShapeMismatch;
}
if (n == 0) {
return Error.DomainError;
}
// The scale every later comparison is relative to, taken once from the input.
var frobenius: CompensatedSum(T) = .{};
for (0..n) |i| {
for (0..n) |j| {
const x: T = try a.at(&.{ i, j });
frobenius.add(x * x);
}
}
const frobenius_norm: T = @sqrt(frobenius.value());
// An all-zero matrix is already diagonal, and its eigenvalues are zero.
const threshold: T = frobenius_norm * @as(T, @floatCast(1.0e-15));
vectors.fill(0);
for (0..n) |i| {
try vectors.setAt(&.{ i, i }, 1);
}
const max_sweeps: usize = 60;
var sweep: usize = 0;
while (sweep < max_sweeps) : (sweep += 1) {
var off: CompensatedSum(T) = .{};
for (0..n) |p| {
for (p + 1..n) |q| {
const x: T = try a.at(&.{ p, q });
off.add(x * x);
}
}
if (@sqrt(off.value()) <= threshold) {
for (0..n) |i| {
values[i] = try a.at(&.{ i, i });
}
return;
}
for (0..n) |p| {
for (p + 1..n) |q| {
const apq: T = try a.at(&.{ p, q });
if (apq == 0) {
continue;
}
// The rotation that zeroes a[p][q]. `t` is the SMALLER root, taken in the form
// that avoids subtracting nearby numbers when theta is large.
const theta: T = ((try a.at(&.{ q, q })) - (try a.at(&.{ p, p }))) / (2 * apq);
const magnitude: T = @abs(theta) + @sqrt(theta * theta + 1);
const t: T = if (theta >= 0) 1 / magnitude else -1 / magnitude;
const c: T = 1 / @sqrt(t * t + 1);
const sn: T = t * c;
// Columns p and q of `a`, then rows p and q, then the accumulated rotations.
for (0..n) |k| {
const kp: T = try a.at(&.{ k, p });
const kq: T = try a.at(&.{ k, q });
try a.setAt(&.{ k, p }, c * kp - sn * kq);
try a.setAt(&.{ k, q }, sn * kp + c * kq);
}
for (0..n) |k| {
const pk: T = try a.at(&.{ p, k });
const qk: T = try a.at(&.{ q, k });
try a.setAt(&.{ p, k }, c * pk - sn * qk);
try a.setAt(&.{ q, k }, sn * pk + c * qk);
}
for (0..n) |k| {
const kp: T = try vectors.at(&.{ k, p });
const kq: T = try vectors.at(&.{ k, q });
try vectors.setAt(&.{ k, p }, c * kp - sn * kq);
try vectors.setAt(&.{ k, q }, sn * kp + c * kq);
}
}
}
}
return Error.DomainError;
}
Cyclic Jacobi: each rotation zeroes one off-diagonal pair, and the sweep repeats until what remains off the diagonal is negligible. Negligible compared to what is the whole question. A fixed threshold declares any sufficiently small matrix already diagonal — scale a matrix by 1e-16 and its off-diagonal sum of squares falls below 1e-30 before a single rotation, so the routine returns the untouched diagonal as the eigenvalues. They are not the eigenvalues, and nothing reports a problem. Measuring against the matrix's own Frobenius norm makes the check scale-free:
| Scale | Residual max |Av − λv| | λ₀ |
|---|---|---|
| 1 | 1.6e-15 | 6.7328e-1 |
| 1e-10 | 1.3e-25 | 6.7328e-11 |
| 1e-20 | 2.9e-35 | 6.7328e-21 |
The test checks A·v = λv for every eigenpair — checking
eigenvalues alone would pass on correct values paired with the wrong vectors — plus the trace
identity as an independent check that none was lost, and orthonormality of the vectors. Running out
of sweeps is a DomainError: Jacobi converges quadratically and reaches working
precision in well under ten sweeps, so returning an unconverged answer silently is the failure
worth avoiding.
try zn.matmulNT(f64, out, a, b); // a · bᵀ
try zn.matmulTN(f64, out, a, b); // aᵀ · b
try zn.solveTriangular(f64, x, r, b, .upper);
try zn.lstsq(f64, x, q, r, a, b); // q and r are scratch and come back as the factorisationmatmulNT/// `out = a * bT`. The right-hand matrix is read transposed rather than materialised.
///
/// THIS IS AN INDEXING CHOICE, NOT AN ALGORITHM. `matmul(out, a, transpose(b))` gives the same
/// answer, but a transposed VIEW makes the inner loop stride across memory; naming the variant
/// lets the loop read both operands the way they are stored. The test asserts it against the
/// materialised form, which is what makes the name safe to reach for.
pub fn matmulNT(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or b.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
if (a.shape[1] != b.shape[1] or out.shape[0] != a.shape[0] or out.shape[1] != b.shape[0]) {
return Error.ShapeMismatch;
}
if (out.isAliased()) {
return Error.UnsupportedShape;
}
for (0..a.shape[0]) |i| {
for (0..b.shape[0]) |j| {
var total: CompensatedSum(T) = .{};
for (0..a.shape[1]) |k| {
total.add((try a.at(&.{ i, k })) * (try b.at(&.{ j, k })));
}
try out.setAt(&.{ i, j }, total.value());
}
}
}
matmulTN/// `out = aT * b`.
pub fn matmulTN(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or b.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
if (a.shape[0] != b.shape[0] or out.shape[0] != a.shape[1] or out.shape[1] != b.shape[1]) {
return Error.ShapeMismatch;
}
if (out.isAliased()) {
return Error.UnsupportedShape;
}
for (0..a.shape[1]) |i| {
for (0..b.shape[1]) |j| {
var total: CompensatedSum(T) = .{};
for (0..a.shape[0]) |k| {
total.add((try a.at(&.{ k, i })) * (try b.at(&.{ k, j })));
}
try out.setAt(&.{ i, j }, total.value());
}
}
}
solveTriangular/// Solve `t * x = b` where `t` is triangular, by substitution. `t` is `(n, n)`, `b` and `x` are
/// `(n, k)`, so several right-hand sides are solved at once.
///
/// `DomainError` on a zero diagonal entry - a triangular matrix with a zero on the diagonal is
/// singular, and substitution would divide by it. Reporting that is more useful than an infinity
/// propagating through the rest of the solve.
pub fn solveTriangular(
comptime T: type,
x: Tensor(T),
t: Tensor(T),
b: Tensor(T),
side: Triangle,
) Error!void {
comptime requireFloat(T);
if (t.rank != 2 or b.rank != 2 or x.rank != 2) {
return Error.UnsupportedShape;
}
const n: usize = t.shape[0];
if (t.shape[1] != n or b.shape[0] != n or x.shape[0] != n or x.shape[1] != b.shape[1]) {
return Error.ShapeMismatch;
}
const columns: usize = b.shape[1];
for (0..columns) |c| {
var step: usize = 0;
while (step < n) : (step += 1) {
// Upper triangles are solved from the bottom row up, lower ones from the top down.
const i: usize = if (side == .upper) n - 1 - step else step;
const pivot: T = try t.at(&.{ i, i });
if (pivot == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
total.add(try b.at(&.{ i, c }));
var j: usize = if (side == .upper) i + 1 else 0;
while (if (side == .upper) j < n else j < i) : (j += 1) {
total.add(-(try t.at(&.{ i, j })) * (try x.at(&.{ j, c })));
}
try x.setAt(&.{ i, c }, total.value() / pivot);
}
}
}
lstsq/// Least squares: the `x` minimising `||a*x - b||`, by QR.
///
/// `a` is `(m, n)` with `m >= n`, `b` is `(m, k)`, `x` is `(n, k)`. `q` `(m, m)` and `r` `(m, n)`
/// are scratch that come back holding the factorisation, and `b` is OVERWRITTEN with `qT*b`.
///
/// QR, NOT THE NORMAL EQUATIONS
///
/// The textbook shortcut solves `aTa * x = aTb`, which is one line and squares the condition
/// number: a system solvable to eight digits becomes solvable to four. Factorising `a = q*r`
/// instead turns the problem into `r*x = qTb` - triangular, and no worse conditioned than `a`.
/// znum does this too and is right to.
///
/// What differs is underneath: this `qr` is Householder where znum's is Gram-Schmidt, so on an
/// ill-conditioned system the `q` here is orthogonal and znum's is not. The test measures the
/// residual on a Hilbert system.
///
/// Nothing is allocated. A caller with several right-hand sides can factor once and call
/// `solveTriangular` directly for each.
pub fn lstsq(
comptime T: type,
x: Tensor(T),
q: Tensor(T),
r: Tensor(T),
a: Tensor(T),
b: Tensor(T),
) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or b.rank != 2 or x.rank != 2) {
return Error.UnsupportedShape;
}
const m: usize = a.shape[0];
const n: usize = a.shape[1];
if (b.shape[0] != m or x.shape[0] != n or x.shape[1] != b.shape[1]) {
return Error.ShapeMismatch;
}
try qr(T, q, r, a);
// b <- qT*b, one column at a time so no second buffer is needed.
var scratch: [max_rank * 512]T = undefined;
if (m > scratch.len) {
return Error.UnsupportedShape;
}
for (0..b.shape[1]) |c| {
for (0..m) |i| {
var total: CompensatedSum(T) = .{};
for (0..m) |k| {
total.add((try q.at(&.{ k, i })) * (try b.at(&.{ k, c })));
}
scratch[i] = total.value();
}
for (0..m) |i| {
try b.setAt(&.{ i, c }, scratch[i]);
}
}
// The leading n rows are the square triangular system; the rest is the residual.
const top_r: Tensor(T) = try r.slice(0, 0, n);
const top_b: Tensor(T) = try b.slice(0, 0, n);
return solveTriangular(T, x, top_r, top_b, .upper);
}
The textbook shortcut for min ‖Ax − b‖ solves
AᵀA·x = Aᵀb. It is one line and it squares the condition
number: a system solvable to eight digits becomes solvable to four. Factorising
A = QR turns the problem into R·x = Qᵀb — triangular,
and no worse conditioned than A. Measured on a Vandermonde design matrix with a known
exact solution:
| n | QR | Normal equations |
|---|---|---|
| 5 | 5.9e-15 | 8.2e-13 |
| 8 | 7.8e-13 | 1.6e-7 |
| 10 | 1.3e-11 | 1.6e-4 |
Both are correct algorithms; one of them stops being usable first. The test asserts the normal equations are visibly worse on the same system, which is what makes the extra factorisation worth its cost rather than a claim in a comment. It also checks the defining property of a least-squares fit — that the residual is orthogonal to every column of the design matrix — on a system with no exact solution.
lstsq allocates nothing: q and r are caller-supplied and
come back holding the factorisation, so a caller with several right-hand sides can factor once and
call solveTriangular for each. b is overwritten with
qᵀb.
matmulNT and matmulTN are an indexing choice rather than an algorithm:
matmul(out, a, transpose(b)) gives the same answer, but a transposed view
makes the inner loop stride across memory. The test asserts both against the materialised form,
which is what makes the names safe to reach for.
try zn.svd(f64, u, &values, v, a); // a = u · diag(values) · vᵀ, values descendingsvd/// Singular value decomposition by one-sided Jacobi: `a = u * diag(values) * vT`.
///
/// `a` is `(m, n)` with `m >= n`. `u` is `(m, n)` and receives the left singular vectors as its
/// columns - it is also the working copy, so `a` is not modified. `values` is `n`, descending.
/// `v` is `(n, n)` with the right singular vectors as its columns.
///
/// ONE-SIDED JACOBI, NOT THE EIGENDECOMPOSITION OF `aTa`
///
/// The short route is `eigh(aTa)` and `sigma = sqrtlambda`, which znum takes. It squares the condition number
/// in exactly the way the normal equations do, and here the consequence is sharper than lost
/// digits: a matrix with singular values 1 and **1e-9** has eigenvalues 1 and **1e-18**, which is
/// below f64's epsilon. The small singular value is not computed inaccurately - **it is gone**,
/// recovered from whatever noise sits in that entry of the Gram matrix.
///
/// One-sided Jacobi rotates PAIRS OF COLUMNS of `a` itself until they are orthogonal. It never
/// forms `aTa`, so every singular value comes back to full relative accuracy however small.
///
/// The test builds a matrix with a known 1e-9 singular value and measures both routes.
///
/// The convergence test is relative, for the same reason as `eigh`: a fixed threshold on the
/// column inner products would call a scaled-down matrix already diagonal.
pub fn svd(
comptime T: type,
u: Tensor(T),
values: []T,
v: Tensor(T),
a: Tensor(T),
) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or u.rank != 2 or v.rank != 2) {
return Error.UnsupportedShape;
}
const m: usize = a.shape[0];
const n: usize = a.shape[1];
if (m < n) {
return Error.UnsupportedShape;
}
if (u.shape[0] != m or u.shape[1] != n or v.shape[0] != n or v.shape[1] != n or
values.len != n)
{
return Error.ShapeMismatch;
}
try materialise(T, u, a);
v.fill(0);
for (0..n) |i| {
try v.setAt(&.{ i, i }, 1);
}
const max_sweeps: usize = 60;
var sweep: usize = 0;
var rotations: usize = 1;
while (sweep < max_sweeps and rotations > 0) : (sweep += 1) {
rotations = 0;
for (0..n) |p| {
for (p + 1..n) |q| {
// The two column lengths and their inner product.
var alpha: CompensatedSum(T) = .{};
var beta: CompensatedSum(T) = .{};
var gamma: CompensatedSum(T) = .{};
for (0..m) |row| {
const ap: T = try u.at(&.{ row, p });
const aq: T = try u.at(&.{ row, q });
alpha.add(ap * ap);
beta.add(aq * aq);
gamma.add(ap * aq);
}
const dot_pq: T = gamma.value();
// Relative: the columns are orthogonal enough when their inner product is
// negligible against their own lengths, at any scale.
const lengths: T = @sqrt(alpha.value() * beta.value());
if (lengths == 0 or @abs(dot_pq) <= lengths * @as(T, @floatCast(1.0e-15))) {
continue;
}
rotations += 1;
const zeta: T = (beta.value() - alpha.value()) / (2 * dot_pq);
const magnitude: T = @abs(zeta) + @sqrt(zeta * zeta + 1);
const t: T = if (zeta >= 0) 1 / magnitude else -1 / magnitude;
const c: T = 1 / @sqrt(t * t + 1);
const sn: T = c * t;
for (0..m) |row| {
const ap: T = try u.at(&.{ row, p });
const aq: T = try u.at(&.{ row, q });
try u.setAt(&.{ row, p }, c * ap - sn * aq);
try u.setAt(&.{ row, q }, sn * ap + c * aq);
}
for (0..n) |row| {
const vp: T = try v.at(&.{ row, p });
const vq: T = try v.at(&.{ row, q });
try v.setAt(&.{ row, p }, c * vp - sn * vq);
try v.setAt(&.{ row, q }, sn * vp + c * vq);
}
}
}
}
if (rotations > 0) {
return Error.DomainError;
}
// Each column's length IS its singular value; normalising leaves the left vectors.
for (0..n) |j| {
var length: CompensatedSum(T) = .{};
for (0..m) |row| {
const x: T = try u.at(&.{ row, j });
length.add(x * x);
}
values[j] = @sqrt(length.value());
if (values[j] > 0) {
for (0..m) |row| {
try u.setAt(&.{ row, j }, (try u.at(&.{ row, j })) / values[j]);
}
}
}
// Descending, by selection - n is small and this keeps the permutation obvious.
for (0..n) |i| {
var best: usize = i;
for (i + 1..n) |j| {
if (values[j] > values[best]) {
best = j;
}
}
if (best == i) {
continue;
}
std.mem.swap(T, &values[i], &values[best]);
for (0..m) |row| {
const keep: T = try u.at(&.{ row, i });
try u.setAt(&.{ row, i }, try u.at(&.{ row, best }));
try u.setAt(&.{ row, best }, keep);
}
for (0..n) |row| {
const keep: T = try v.at(&.{ row, i });
try v.setAt(&.{ row, i }, try v.at(&.{ row, best }));
try v.setAt(&.{ row, best }, keep);
}
}
}
The short route is eigh(aᵀa) with σ = √λ. It
squares the condition number, and here the consequence is sharper than lost digits: a matrix with
singular values 1 and 1e-9 has eigenvalues 1 and 1e-18, which is
below f64's epsilon. The small singular value is not computed inaccurately — it is
gone, recovered from whatever noise sits in that entry of the Gram matrix.
One-sided Jacobi rotates pairs of columns of a itself until they are
orthogonal, never forming aᵀa. Relative error in the smallest singular value:
| True σmin | One-sided Jacobi | Through aᵀa |
|---|---|---|
| 1e-3 | 1.5e-14 | 1.7e-11 |
| 1e-6 | 4.6e-12 | 1.1e-5 |
| 1e-9 | 2.0e-8 | 5.7 — 570% wrong |
At 1e-9 the Gram route returns 6.7e-9 for a value of 1e-9. That is not a precision difference; it is a different number. The test asserts the gap only where the squaring actually bites — at 1e-3 both routes are fine, and claiming otherwise would overstate the case.
Reconstruction is checked too: u · diag(σ) · vᵀ must be
a again, because checking the singular values alone would pass on correct values
paired with the wrong vectors. Each column's length is its singular value, so normalising
the working copy at the end leaves the left vectors in place with no extra pass.
try zn.pinv(f64, out, u, &values, v, a, 1.0e-15); // u, values, v are scratch
try zn.inverse(f64, out, a, pivots); // a is overwrittenpinv/// The Moore-Penrose pseudo-inverse, via `svd`: `out = v * diag(1/sigma) * uT`, with the singular
/// values below `rcond * sigma_max` treated as zero.
///
/// `a` is `(m, n)` with `m >= n`; `out` is `(n, m)`. `u` `(m, n)`, `v` `(n, n)` and `values` `n`
/// are scratch and come back holding the decomposition.
///
/// THE RANK CUTOFF IS RELATIVE TO THE LARGEST SINGULAR VALUE
///
/// Deciding which singular values are "zero" is the whole content of a pseudo-inverse - it is
/// what makes it defined for a singular matrix at all. An ABSOLUTE cutoff, which znum uses at
/// `1e-12`, gets this wrong in both directions: **a matrix scaled down by 1e-6 has every
/// singular value discarded and the pseudo-inverse comes back as zeros**, while one scaled up by
/// 1e12 keeps directions that are numerically noise and inverts them into enormous numbers.
///
/// `rcond * sigma_max` is the standard, and it is scale-free by construction: multiplying `a` by
/// any constant multiplies `sigma_max` by the same one, so exactly the same directions are kept.
/// The test asserts that - `pinv(k*a) = pinv(a)/k` across twelve orders of magnitude.
///
/// The default `rcond` is `1e-15`, near f64's epsilon: a direction whose singular value is
/// that far below the largest carries no information the input's own precision can support.
pub fn pinv(
comptime T: type,
out: Tensor(T),
u: Tensor(T),
values: []T,
v: Tensor(T),
a: Tensor(T),
rcond: T,
) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
const m: usize = a.shape[0];
const n: usize = a.shape[1];
if (out.shape[0] != n or out.shape[1] != m) {
return Error.ShapeMismatch;
}
if (rcond < 0) {
return Error.DomainError;
}
try svd(T, u, values, v, a);
// `values` is descending, so the first is the largest.
const cutoff: T = values[0] * rcond;
for (0..n) |i| {
for (0..m) |j| {
var total: CompensatedSum(T) = .{};
for (0..n) |k| {
if (values[k] <= cutoff) {
continue;
}
total.add((try v.at(&.{ i, k })) * (try u.at(&.{ j, k })) / values[k]);
}
try out.setAt(&.{ i, j }, total.value());
}
}
}
inverse/// The inverse of a square matrix, by solving `a * x = I`.
///
/// `a` and `work` are `(n, n)`; both are overwritten. `pivots` needs `n` entries.
///
/// Solving against the identity rather than computing a closed form: `solve` already does the
/// factorisation with partial pivoting, and an inverse is `n` right-hand sides of the same
/// system. **Wanting an inverse is usually a sign of wanting a solve** - `solve(a, b)` is more
/// accurate and cheaper than `matmul(inverse(a), b)`, and this exists for the cases where the
/// matrix itself is the answer.
pub fn inverse(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
pivots: []usize,
) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
const n: usize = a.shape[0];
if (a.shape[1] != n or out.shape[0] != n or out.shape[1] != n) {
return Error.ShapeMismatch;
}
out.fill(0);
for (0..n) |i| {
try out.setAt(&.{ i, i }, 1);
}
return solve(T, a, out, pivots);
}
Deciding which singular values count as zero is the content of a pseudo-inverse — it is what makes one defined for a singular matrix at all. An absolute cutoff gets this wrong in both directions. A matrix with singular values 1, 1e-3, 1e-7, scaled:
| Scale | Singular values | Absolute cutoff at 1e-12 |
|---|---|---|
| 1 | 1e0, 1e-3, 1e-7 | keeps 3 of 3 |
| 1e-6 | 1e-6, 1e-9, 1e-13 | keeps 2 of 3 — a real direction, dropped |
Nothing about the matrix changed except its units. rcond · σmax
is scale-free by construction: multiplying a by any constant multiplies
σmax by the same one, so exactly the same directions are kept. The
test asserts pinv(k·a) = pinv(a)/k to nine digits across six orders of
magnitude, and checks the defining Penrose condition a · a⁺ · a =
a — which holds for a rank-deficient matrix too, and is the reason the pseudo-inverse
is worth having.
inverse solves against the identity rather than computing a closed form.
Wanting an inverse is usually a sign of wanting a solve: solve(a, b)
is both more accurate and cheaper than matmul(inverse(a), b), and this exists for the
cases where the matrix itself is the answer.
try zn.cholesky(f64, lower, a); // lower * lowerT == acholesky/// The Cholesky factor of a symmetric positive-definite matrix: lower-triangular `L` with
/// `L LT = a`. Written into `out`, which is zeroed above the diagonal.
///
/// A NON-POSITIVE PIVOT IS THE DEFINITION OF NOT POSITIVE-DEFINITE, so `DomainError` here is
/// informative rather than a limitation: attempting the factorisation IS the standard test for
/// the property, and it is cheaper than computing eigenvalues.
pub fn cholesky(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or a.shape[0] != a.shape[1]) {
return Error.UnsupportedShape;
}
if (out.rank != 2 or out.shape[0] != a.shape[0] or out.shape[1] != a.shape[1]) {
return Error.ShapeMismatch;
}
const n: usize = a.shape[0];
out.fill(0);
var i: usize = 0;
while (i < n) : (i += 1) {
var j: usize = 0;
while (j <= i) : (j += 1) {
var acc: T = a.at2(i, j);
var k: usize = 0;
while (k < j) : (k += 1) {
acc -= out.at2(i, k) * out.at2(j, k);
}
if (i == j) {
if (acc <= 0) {
return Error.DomainError;
}
out.setAt2(i, j, @sqrt(acc));
} else {
const diag: T = out.at2(j, j);
out.setAt2(i, j, acc / diag);
}
}
}
}
For a symmetric positive-definite a, produces lower-triangular L with
L Lᵀ = a. Unlike the LU routines this writes to a separate
output, which it zeroes above the diagonal.
A non-positive value on the diagonal during the factorisation means the matrix is not
positive-definite, and the call returns DomainError. Attempting the factorisation is
the standard test for the property and is cheaper than computing eigenvalues, so the error is
informative rather than a limitation.
try zn.tensordot(f64, out, a, b, axis_a, axis_b); // the general contraction
try zn.bmm(f64, out, batched_a, batched_b);
try zn.outer(f64, out, v1, v2);
try zn.diagonal(f64, diag, m);
try zn.triangle(f64, out, m, .upper, true);
const cond = try zn.conditionNumber(f64, u, &sigma, v, a);tensordot/// Contract `a` and `b` over one axis each: sum over `a`'s `axis_a` and `b`'s `axis_b`, which
/// must have the same extent. The output's axes are `a`'s remaining ones followed by `b`'s.
///
/// THE GENERAL CASE OF EVERY PRODUCT IN THIS FILE. `matmul` is `tensordot` over `a`'s last axis
/// and `b`'s first; `dotAll` on rank-1 inputs is `tensordot` with nothing left over; an outer
/// product is `tensordot` over no axes at all. Having the general one written once means the
/// special ones can be checked against it, which the test does.
///
/// Slower than `matmul` for the rank-2 case, by roughly the cost of the index arithmetic. The
/// specialised versions stay for that reason and are what a hot loop should call.
pub fn tensordot(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
b: Tensor(T),
axis_a: usize,
axis_b: usize,
) Error!void {
comptime requireFloat(T);
if (axis_a >= a.rank or axis_b >= b.rank) {
return Error.OutOfRange;
}
const contracted: usize = a.shape[axis_a];
if (b.shape[axis_b] != contracted) {
return Error.ShapeMismatch;
}
if (out.rank != a.rank + b.rank - 2 or out.isAliased()) {
return Error.ShapeMismatch;
}
// The output's axes are a's surviving ones then b's, in order.
var free_a: [max_rank]usize = @splat(0);
var free_b: [max_rank]usize = @splat(0);
var count_a: usize = 0;
var count_b: usize = 0;
for (0..a.rank) |i| {
if (i != axis_a) {
free_a[count_a] = i;
count_a += 1;
}
}
for (0..b.rank) |i| {
if (i != axis_b) {
free_b[count_b] = i;
count_b += 1;
}
}
for (0..count_a) |i| {
if (out.shape[i] != a.shape[free_a[i]]) {
return Error.ShapeMismatch;
}
}
for (0..count_b) |i| {
if (out.shape[count_a + i] != b.shape[free_b[i]]) {
return Error.ShapeMismatch;
}
}
var walker: [max_rank]usize = @splat(0);
var index_a: [max_rank]usize = @splat(0);
var index_b: [max_rank]usize = @splat(0);
var remaining: usize = out.size();
while (remaining > 0) : (remaining -= 1) {
for (0..count_a) |i| {
index_a[free_a[i]] = walker[i];
}
for (0..count_b) |i| {
index_b[free_b[i]] = walker[count_a + i];
}
var total: CompensatedSum(T) = .{};
for (0..contracted) |k| {
index_a[axis_a] = k;
index_b[axis_b] = k;
total.add((try a.at(index_a[0..a.rank])) * (try b.at(index_b[0..b.rank])));
}
try out.setAt(walker[0..out.rank], total.value());
advance(out.shape[0..out.rank], walker[0..out.rank]);
}
}
bmm/// Batched matrix product: `out[b] = a[b] @ b_operand[b]` for every index of the leading axis.
///
/// Rank 3 throughout, `(batch, m, k) @ (batch, k, n) -> (batch, m, n)`. A loop over `matmul` on
/// slices, which is what it is - the value is that the shape check happens once and covers every
/// batch, rather than the caller writing the loop and checking nothing.
pub fn bmm(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 3 or b.rank != 3 or out.rank != 3) {
return Error.UnsupportedShape;
}
const batch: usize = a.shape[0];
if (b.shape[0] != batch or out.shape[0] != batch) {
return Error.ShapeMismatch;
}
if (a.shape[2] != b.shape[1] or out.shape[1] != a.shape[1] or out.shape[2] != b.shape[2]) {
return Error.ShapeMismatch;
}
for (0..batch) |i| {
const left: Tensor(T) = try (try a.slice(0, i, 1)).squeeze(0);
const right: Tensor(T) = try (try b.slice(0, i, 1)).squeeze(0);
const result: Tensor(T) = try (try out.slice(0, i, 1)).squeeze(0);
try matmul(T, result, left, right);
}
}
outer/// The outer product `out[i][j] = a[i] * b[j]` of two rank-1 tensors.
pub fn outer(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
comptime requireNumeric(T);
if (a.rank != 1 or b.rank != 1 or out.rank != 2) {
return Error.UnsupportedShape;
}
if (out.shape[0] != a.shape[0] or out.shape[1] != b.shape[0]) {
return Error.ShapeMismatch;
}
for (0..a.shape[0]) |i| {
const x: T = try a.at(&.{i});
for (0..b.shape[0]) |j| {
try out.setAt(&.{ i, j }, x * (try b.at(&.{j})));
}
}
}
diagonal/// The diagonal of a rank-2 tensor, into a rank-1 `out`. Not square-only: the diagonal of a
/// rectangular matrix is its `min(rows, cols)` leading entries.
pub fn diagonal(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireNumeric(T);
if (a.rank != 2 or out.rank != 1) {
return Error.UnsupportedShape;
}
const n: usize = @min(a.shape[0], a.shape[1]);
if (out.shape[0] != n) {
return Error.ShapeMismatch;
}
for (0..n) |i| {
try out.setAt(&.{i}, try a.at(&.{ i, i }));
}
}
triangle/// Copy one triangle of `a` into `out` and zero the other. `keep_diagonal` decides whether the
/// diagonal itself belongs to the kept side.
pub fn triangle(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
side: Triangle,
keep_diagonal: bool,
) Error!void {
comptime requireNumeric(T);
if (a.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
if (out.shape[0] != a.shape[0] or out.shape[1] != a.shape[1] or out.isAliased()) {
return Error.ShapeMismatch;
}
for (0..a.shape[0]) |i| {
for (0..a.shape[1]) |j| {
const on_diagonal: bool = i == j;
const above: bool = j > i;
const keep: bool = if (on_diagonal)
keep_diagonal
else if (side == .upper)
above
else
!above;
try out.setAt(&.{ i, j }, if (keep) try a.at(&.{ i, j }) else 0);
}
}
}
conditionNumber/// The 2-norm condition number: `sigma_max / sigma_min`, via `svd`.
///
/// THE NUMBER THAT EXPLAINS THE REST OF THIS FILE
///
/// It says how much a relative error in the input can be multiplied on its way to the output. A
/// system with condition 1e8 solved in f64 - sixteen digits - keeps about eight. That single
/// quantity is behind every algorithm choice recorded here: why `qr` is Householder, why `lstsq`
/// avoids the normal equations (which SQUARE this number), why `svd` does not go through the
/// Gram matrix, and why `pinv`'s cutoff is relative.
///
/// Infinite for a singular matrix, which is the honest answer rather than a large finite one:
/// the smallest singular value is zero and the ratio is unbounded.
///
/// `u`, `values` and `v` are scratch, sized as `svd` requires, and come back holding the
/// decomposition. `a` is overwritten.
pub fn conditionNumber(
comptime T: type,
u: Tensor(T),
values: []T,
v: Tensor(T),
a: Tensor(T),
) Error!T {
comptime requireFloat(T);
try svd(T, u, values, v, a);
const smallest: T = values[values.len - 1];
if (smallest == 0) {
return inf(T);
}
return values[0] / smallest;
}
tensordot is the general case of every product in this file. matmul
is a contraction over a's last axis and b's first;
matmulNT is over both last axes; a rank-1 contraction with nothing left over is the
dot product; an outer product is a contraction over nothing. So this is exact for
each specialisation against the general form rather than against hand-written
numbers — it is the specialisations that are likely to be wrong and the general form that is
easy to reason about.
conditionNumber is σmax / σmin, and
it is the number that explains the rest of this file. It says how much a relative
error in the input can be multiplied on the way to the output: a system with condition 1e8 solved
in f64 — sixteen digits — keeps about eight. That single quantity is behind every
algorithm choice recorded here — why qr is Householder, why lstsq
avoids the normal equations (which square it), why svd does not go through
the Gram matrix, why pinv's cutoff is relative.
It is infinite for a singular matrix, which is the honest answer rather than a large finite one. A 5×5 Hilbert matrix at roughly 4.8e5, which is why that matrix is the one every algorithm here is tested on.
try zn.einsum(f64, "ij,jk->ik", out, .{ a, b }); // matrix product
try zn.einsum(f64, "ij->ji", out, .{a}); // transpose
try zn.einsum(f64, "ii->", scalar, .{m}); // trace
try zn.einsum(f64, "ii->i", diag, .{m}); // diagonal
try zn.einsum(f64, "i,j->ij", out, .{ v1, v2 }); // outer product
try zn.einsum(f64, "bij,bjk->bik", out, .{ ba, bb }); // batchedeinsum/// Einstein summation: `einsum(f64, "ij,jk->ik", out, .{ a, b })` is a matrix product.
///
/// Every label names an axis. A label in the output is a free axis; a label appearing only in the
/// inputs is summed over. That one rule covers matrix products, transposes, traces, outer
/// products, batched contractions and diagonal extraction, which is why the notation exists.
///
/// The spec is comptime - see `parseEinsum` for what that buys. `operands` is a tuple, so a
/// one-operand form (`"ii->"` for a trace, `"ij->ji"` for a transpose) is written the same way as
/// a two-operand one.
///
/// A REPEATED LABEL WITHIN ONE OPERAND SELECTS ITS DIAGONAL, which is how `"ii->"` is a trace.
/// The index assignment sets both axes to the same value, so the walk visits only the diagonal
/// without anything special being written for it.
pub fn einsum(
comptime T: type,
comptime spec: []const u8,
out: Tensor(T),
operands: anytype,
) Error!void {
comptime requireFloat(T);
const parsed: EinsumSpec = comptime parseEinsum(spec, operands.len);
if (out.rank != parsed.output.len or out.isAliased()) {
return Error.ShapeMismatch;
}
// Each label's extent, taken from wherever it first appears and checked everywhere else.
var extent: [256]usize = @splat(0);
var known: [256]bool = @splat(false);
inline for (0..parsed.input_count) |k| {
const operand: Tensor(T) = operands[k];
if (operand.rank != parsed.inputs[k].len) {
return Error.ShapeMismatch;
}
for (parsed.inputs[k], 0..) |label, axis| {
if (known[label]) {
if (extent[label] != operand.shape[axis]) {
return Error.ShapeMismatch;
}
} else {
extent[label] = operand.shape[axis];
known[label] = true;
}
}
}
for (parsed.output, 0..) |label, axis| {
if (out.shape[axis] != extent[label]) {
return Error.ShapeMismatch;
}
}
// The contracted axes, as a shape to walk.
var sum_shape: [max_labels]usize = @splat(1);
for (parsed.summed, 0..) |label, i| {
sum_shape[i] = extent[label];
}
var sum_total: usize = 1;
for (0..parsed.summed.len) |i| {
sum_total *= sum_shape[i];
}
var position: [256]usize = @splat(0);
var out_walker: [max_rank]usize = @splat(0);
var index: [max_rank]usize = @splat(0);
var remaining_out: usize = out.size();
while (remaining_out > 0) : (remaining_out -= 1) {
for (parsed.output, 0..) |label, axis| {
position[label] = out_walker[axis];
}
var total: CompensatedSum(T) = .{};
var sum_walker: [max_labels]usize = @splat(0);
var remaining_sum: usize = sum_total;
while (remaining_sum > 0) : (remaining_sum -= 1) {
for (parsed.summed, 0..) |label, i| {
position[label] = sum_walker[i];
}
var term: T = 1;
inline for (0..parsed.input_count) |k| {
const operand: Tensor(T) = operands[k];
for (parsed.inputs[k], 0..) |label, axis| {
index[axis] = position[label];
}
term *= try operand.at(index[0..parsed.inputs[k].len]);
}
total.add(term);
advance(sum_shape[0..parsed.summed.len], sum_walker[0..parsed.summed.len]);
}
try out.setAt(out_walker[0..out.rank], total.value());
advance(out.shape[0..out.rank], out_walker[0..out.rank]);
}
}
Every label names an axis. A label in the output is a free axis; a label appearing only in the
inputs is summed over. That one rule covers all six lines above — and a repeated label
within one operand selects its diagonal, which is how "ii->" is a trace with
nothing written for the diagonal case at all.
The spec is comptime, and that is the divergence. znum takes
spec: []const u8 and parses it at runtime, so a typo, an unknown output label or too
many operands all come back as an Error the caller must handle — at the point of
use, on data. But every one of those is a property of the spec, which is written in the
source and cannot change while the program runs.
| Written | Result |
|---|---|
"ij,jk" | compile error: spec needs an explicit -> |
"ij,jk->iq" | compile error: output label q appears in no input |
"ij,jk->ii" | compile error: output label i is repeated |
"ij->ij" with two operands | compile error: spec names 1 operand, 2 given |
What is left for runtime is what genuinely depends on the data: a rank that does not match its
subscript, and two operands disagreeing about a label's extent. Those stay
ShapeMismatch.
The test checks each form against the function it duplicates — matmul,
trace, diagonal, outer, dotAll,
sumAxis, bmm — rather than against a table of expected numbers. A
table would test einsum against whoever wrote the table.
try zn.softmaxRows(f32, probs, logits); // rank 2; each ROW becomes a distributionsoftmaxRows/// Softmax along the last axis of a rank-2 tensor: each ROW becomes a probability distribution.
///
/// THE ROW MAXIMUM IS SUBTRACTED, AND IT IS NOT AN OPTIMISATION
///
/// `e^x` overflows `f32` above about 88, and logits reach that routinely in an untrained network.
/// Subtracting the row's maximum first leaves the result identical - softmax is invariant under
/// adding a constant to a whole row, which is asserted in the tests - while bounding the largest
/// exponent at `e^0 = 1`. Without it the first forward pass of a badly initialised layer returns
/// NaN, and the failure looks like a bad initialisation rather than a bad softmax.
///
/// Rank 2 only, and the rows are the last axis. A caller with a different layout transposes,
/// which is a view and costs nothing.
pub fn softmaxRows(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
if (out.shape[0] != a.shape[0] or out.shape[1] != a.shape[1] or out.isAliased()) {
return Error.ShapeMismatch;
}
const rows: usize = a.shape[0];
const cols: usize = a.shape[1];
if (cols == 0) {
return Error.DomainError;
}
var row: usize = 0;
while (row < rows) : (row += 1) {
var largest: T = a.at2(row, 0);
var col: usize = 1;
while (col < cols) : (col += 1) {
const x: T = a.at2(row, col);
if (x > largest) {
largest = x;
}
}
// Compensated, like every other accumulation here: a wide row is exactly the sum that
// loses its low bits, and the result is a denominator every element is divided by.
var total: CompensatedSum(T) = .{};
col = 0;
while (col < cols) : (col += 1) {
const e: T = @exp(a.at2(row, col) - largest);
out.setAt2(row, col, e);
total.add(e);
}
const denominator: T = total.value();
col = 0;
while (col < cols) : (col += 1) {
const e: T = out.at2(row, col);
out.setAt2(row, col, e / denominator);
}
}
}
The row maximum is subtracted before exponentiating. e^x overflows
f32 above about 88 and logits reach that in an untrained network; subtracting the
maximum leaves the result identical, because softmax is invariant under adding a constant to a
whole row, while bounding the largest exponent at e0. The row sum is
compensated, since every element of the row is divided by it.
try zn.layerNormRows(f32, out, a, 1.0e-5); // epsilon is yours to passlayerNormRows/// Normalise each row of a rank-2 tensor to zero mean and unit variance.
///
/// `out[r][c] = (a[r][c] - mean_r) / sqrt(var_r + epsilon)`, with the POPULATION variance.
///
/// WHY EPSILON IS A PARAMETER AND NOT A CONSTANT
///
/// A row of identical values has zero variance, and dividing by its square root is an infinity
/// that propagates through the rest of the network. Every implementation adds a small epsilon
/// inside the root; the VALUE differs between frameworks (PyTorch defaults to 1e-5, JAX to 1e-6)
/// and it changes results measurably in `f32`. Baking one in would make this library silently
/// disagree with whichever the caller expected, so it is theirs to pass.
///
/// Population variance, not sample: normalisation is a statement about the row in hand, not an
/// estimate of a population it was drawn from. `Ddof` exists for the case where that distinction
/// matters and this is not it.
pub fn layerNormRows(
comptime T: type,
out: Tensor(T),
a: Tensor(T),
epsilon: T,
) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
if (out.shape[0] != a.shape[0] or out.shape[1] != a.shape[1] or out.isAliased()) {
return Error.ShapeMismatch;
}
const rows: usize = a.shape[0];
const cols: usize = a.shape[1];
if (cols == 0) {
return Error.DomainError;
}
const width: T = @floatFromInt(cols);
var row: usize = 0;
while (row < rows) : (row += 1) {
// Both accumulations are compensated, for the reason recorded with `sumAll`: a wide row
// is exactly the sum that loses its low bits, and here it feeds a DIVISOR.
var total: CompensatedSum(T) = .{};
var col: usize = 0;
while (col < cols) : (col += 1) {
const x: T = a.at2(row, col);
total.add(x);
}
const mean: T = total.value() / width;
var sq: CompensatedSum(T) = .{};
col = 0;
while (col < cols) : (col += 1) {
const d: T = a.at2(row, col) - mean;
const term: T = d * d;
sq.add(term);
}
const inv: T = 1.0 / @sqrt(sq.value() / width + epsilon);
col = 0;
while (col < cols) : (col += 1) {
const x: T = a.at2(row, col);
out.setAt2(row, col, (x - mean) * inv);
}
}
}
Each row comes out with mean 0 and population variance 1. Three passes — mean, variance,
write. The single-pass identity E[x²] − E[x]² is not used: it subtracts two large
nearly-equal numbers and loses most of its significant digits when the mean is large relative to
the spread, which is the normal case for an unnormalised activation.
Epsilon is a parameter because frameworks disagree about its value — PyTorch defaults to
1e-5, JAX to 1e-6 — and the difference is visible in f32. A row of identical
values has zero variance, and epsilon is what keeps its result finite.
const loss: f32 = try zn.mseLoss(f32, prediction, target);
try zn.sgdStep(f32, weights, weights, gradient, 0.1); // in placemseLoss/// Mean squared error between two tensors of the same shape.
///
/// Returns a scalar rather than writing into a tensor, because a loss IS a scalar and handing
/// back a rank-0 tensor would make every caller unwrap it. `DomainError` for empty input, since
/// the mean of nothing is a question for the caller.
pub fn mseLoss(comptime T: type, prediction: Tensor(T), target: Tensor(T)) Error!T {
comptime requireFloat(T);
if (prediction.rank != target.rank) {
return Error.ShapeMismatch;
}
var axis: usize = 0;
while (axis < prediction.rank) : (axis += 1) {
if (prediction.shape[axis] != target.shape[axis]) {
return Error.ShapeMismatch;
}
}
const n: usize = prediction.size();
if (n == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var walk: Walk = .over(prediction.shape[0..prediction.rank]);
while (walk.next()) |at| {
const delta: T = (try prediction.at(at)) - (try target.at(at));
const squared: T = delta * delta;
total.add(squared);
}
const divisor: T = @floatFromInt(n);
return total.value() / divisor;
}
sgdStep/// `out = weight - rate * grad` - one step of plain gradient descent.
///
/// NOT IN PLACE, and the caller may still pass `weight` as `out` to get in-place behaviour.
/// Returning a fresh result by default means the operation composes and can be tested against a
/// reference without the reference having been consumed; an in-place-only signature forces every
/// caller to copy first, including the tests.
///
/// Written as a loop rather than through `zip`, because `zip`'s function is a COMPTIME
/// parameter and cannot close over a runtime `rate`. A `zipScalar` seam would fix that and is
/// not worth inventing for one caller.
pub fn sgdStep(
comptime T: type,
out: Tensor(T),
weight: Tensor(T),
grad: Tensor(T),
rate: T,
) Error!void {
comptime requireFloat(T);
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const vw: Tensor(T) = try weight.broadcastTo(shape);
const vg: Tensor(T) = try grad.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const w: T = vw.data[vw.offsetOf(at)];
const g: T = vg.data[vg.offsetOf(at)];
out.data[out.offsetOf(at)] = w - rate * g;
}
}
mseLoss returns a scalar rather than a rank-0 tensor, and does not broadcast:
shapes must match exactly. sgdStep computes weight - rate * grad into a
separate destination by default; passing weight as out gives the
in-place form, and the two agree.
const mse = try zn.mseLoss(f32, prediction, target);
const mae = try zn.maeLoss(f32, prediction, target);
const huber = try zn.huberLoss(f32, prediction, target, 1.0);
const bce = try zn.binaryCrossEntropyFromLogits(f32, logits, target);
const kl = try zn.klDivergenceRows(f32, p, q);mseLoss/// Mean squared error between two tensors of the same shape.
///
/// Returns a scalar rather than writing into a tensor, because a loss IS a scalar and handing
/// back a rank-0 tensor would make every caller unwrap it. `DomainError` for empty input, since
/// the mean of nothing is a question for the caller.
pub fn mseLoss(comptime T: type, prediction: Tensor(T), target: Tensor(T)) Error!T {
comptime requireFloat(T);
if (prediction.rank != target.rank) {
return Error.ShapeMismatch;
}
var axis: usize = 0;
while (axis < prediction.rank) : (axis += 1) {
if (prediction.shape[axis] != target.shape[axis]) {
return Error.ShapeMismatch;
}
}
const n: usize = prediction.size();
if (n == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var walk: Walk = .over(prediction.shape[0..prediction.rank]);
while (walk.next()) |at| {
const delta: T = (try prediction.at(at)) - (try target.at(at));
const squared: T = delta * delta;
total.add(squared);
}
const divisor: T = @floatFromInt(n);
return total.value() / divisor;
}
maeLoss/// Mean absolute error between two tensors of the same shape.
///
/// Not differentiable at zero, which is the whole trade against `mseLoss`: it does not square,
/// so an outlier contributes its distance rather than its distance squared, and a single bad
/// sample cannot dominate the batch.
pub fn maeLoss(comptime T: type, a: Tensor(T), b: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n != b.size() or a.rank != b.rank) {
return Error.ShapeMismatch;
}
if (n == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var pair: Walk2 = .over(a.shape[0..a.rank], b.shape[0..b.rank]);
while (pair.next()) |at| {
total.add(@abs((try a.at(at.first)) - (try b.at(at.second))));
}
return total.value() / @as(T, @floatFromInt(n));
}
huberLoss/// Huber loss: squared below `delta`, linear above it.
///
/// THE TWO BRANCHES MEET, AND SO DO THEIR SLOPES
///
/// At `|d| = delta` the quadratic gives `delta^2/2` and the linear branch gives
/// `delta*delta - delta^2/2 = delta^2/2`. Their derivatives are `delta` on both sides. That
/// smoothness is the point: Huber is MSE where the error is small and MAE where it is large, and
/// a join with a kink would put a discontinuity in the gradient exactly at the threshold most
/// samples sit near.
///
/// `delta` must be positive; zero would make the loss identically linear and the parameter
/// meaningless.
pub fn huberLoss(comptime T: type, a: Tensor(T), b: Tensor(T), delta: T) Error!T {
comptime requireFloat(T);
if (delta <= 0) {
return Error.DomainError;
}
const n: usize = a.size();
if (n != b.size() or a.rank != b.rank) {
return Error.ShapeMismatch;
}
if (n == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var pair: Walk2 = .over(a.shape[0..a.rank], b.shape[0..b.rank]);
while (pair.next()) |at| {
const d: T = @abs((try a.at(at.first)) - (try b.at(at.second)));
total.add(if (d <= delta) 0.5 * d * d else delta * (d - 0.5 * delta));
}
return total.value() / @as(T, @floatFromInt(n));
}
binaryCrossEntropyFromLogits/// Binary cross-entropy from LOGITS, averaged over elements. `targets` are in [0, 1].
///
/// COMPUTED AS `max(x,0) - x*t + log(1 + e^-|x|)`
///
/// The definition applies a sigmoid and then a logarithm, and both ends break: a logit of -800
/// gives a sigmoid of exactly 0 and a logarithm of `-inf`, and a logit of +800 gives 1 and the
/// same failure on the other side. The identity above is algebraically equal for every real `x`
/// and contains only `e^-|x|`, whose argument is never positive, so it cannot overflow.
///
/// Taking logits rather than probabilities is not a convenience. A caller who has already
/// applied a sigmoid has lost the information this needs, and no implementation can recover it.
pub fn binaryCrossEntropyFromLogits(
comptime T: type,
logits: Tensor(T),
targets: Tensor(T),
) Error!T {
comptime requireFloat(T);
const n: usize = logits.size();
if (n != targets.size() or logits.rank != targets.rank) {
return Error.ShapeMismatch;
}
if (n == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var pair: Walk2 = .over(logits.shape[0..logits.rank], targets.shape[0..targets.rank]);
while (pair.next()) |at| {
const x: T = try logits.at(at.first);
const t: T = try targets.at(at.second);
total.add(@max(x, 0) - x * t + @log(1 + @exp(-@abs(x))));
}
return total.value() / @as(T, @floatFromInt(n));
}
klDivergenceRows/// Kullback-Leibler divergence from `q` to `p`, row by row, averaged over rows. Both are rank-2
/// and each row must be a distribution.
///
/// NOT SYMMETRIC, and the argument order says which way round it goes: this is
/// `sum p*log(p/q)`, the divergence OF `q` FROM `p`, which is the one a loss wants - `p` is the
/// target and the penalty grows without bound where `q` assigns no probability to something `p`
/// does.
///
/// Terms where `p` is zero contribute nothing, by the limit of `p*log p`. A zero in `q` where
/// `p` is not zero is genuinely infinite, and `DomainError` says so rather than returning `inf`.
pub fn klDivergenceRows(comptime T: type, p: Tensor(T), q: Tensor(T)) Error!T {
comptime requireFloat(T);
if (p.rank != 2 or q.rank != 2) {
return Error.UnsupportedShape;
}
if (p.shape[0] != q.shape[0] or p.shape[1] != q.shape[1]) {
return Error.ShapeMismatch;
}
const rows: usize = p.shape[0];
if (rows == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var row: usize = 0;
while (row < rows) : (row += 1) {
var col: usize = 0;
while (col < p.shape[1]) : (col += 1) {
const pv: T = try p.at(&.{ row, col });
const qv: T = try q.at(&.{ row, col });
if (pv < 0 or qv < 0) {
return Error.DomainError;
}
if (pv == 0) {
continue;
}
if (qv == 0) {
return Error.DomainError;
}
total.add(pv * @log(pv / qv));
}
}
return total.value() / @as(T, @floatFromInt(rows));
}
All return a scalar and require matching sizes; none broadcasts.
maeLoss does not square, so an outlier contributes its distance rather than its
distance squared and cannot dominate a batch. It is not differentiable at zero.
huberLoss is squared below delta and linear above it:
total += if (d <= delta) 0.5 * d * d else delta * (d - 0.5 * delta);At |d| = delta both branches give delta²/2 and both derivatives
give delta. That smoothness is the point: a kink at the join would put a
discontinuity in the gradient exactly at the threshold most samples sit near. delta
must be positive.
binaryCrossEntropyFromLogits takes logits, for cross-entropy's reason. The
definition applies a sigmoid then a logarithm, and both ends break: a logit of −800 gives a
sigmoid of exactly 0 and a logarithm of −inf, and +800 fails the same way on the other side.
The form used is
max(x, 0) − x·t + log(1 + e−|x|), algebraically equal for
every real x and containing only a non-positive exponent, so it cannot overflow. On
logits of ±800 it returns exactly 0 for a correct target and exactly 800 for a wrong one.
klDivergenceRows is ∑ p·log(p/q) per row, averaged — the
divergence of q from p, which is the direction a loss
wants, with p the target. It is not symmetric, so the argument order matters. Terms
where p is zero contribute nothing; a zero in q where p is
not is genuinely infinite and returns DomainError rather than an infinity.
try zn.conv2d(f64, out, image, kernel, .{ 1, 1 }, .{ 0, 0 }); // stride, pad per axis
try zn.conv2d(f64, out, image, kernel, .{ 1, 1 }, .{ 1, 1 }); // 3x3 kernel: "same" size
try zn.conv2d(f64, out, image, kernel, .{ 2, 2 }, .{ 0, 0 }); // half-size output
try zn.maxPool2d(f64, out, image, 2); // (h/2, w/2)
try zn.avgPool2d(f64, out, image, 2);conv2d/// 2-D cross-correlation of a single-channel image with a kernel, with stride and zero padding.
/// The output is `((h + 2*pad - kh) / stride + 1, (w + 2*pad - kw) / stride + 1)`.
///
/// CROSS-CORRELATION, NOT CONVOLUTION, AND THE NAME SAYS CONVOLUTION ANYWAY
///
/// Every deep-learning framework's "conv2d" slides the kernel WITHOUT flipping it, which is
/// cross-correlation. A true convolution flips the kernel first. Since the kernel is learned, the
/// distinction is invisible to training - a flipped kernel is just a different kernel - and the
/// unflipped form is one indexing operation cheaper. The name follows every framework rather
/// than the mathematics, because that is what a reader will look for.
///
/// `stride` and `pad` are separate per axis, matching znum's signature exactly. Padding is
/// implicit: a tap that lands outside the image contributes zero, and no padded copy is made.
/// `pad = kh / 2` with `stride = 1` gives an output the size of the input.
///
/// Rank 2, one channel. The batched multi-channel form is a loop over this one.
pub fn conv2d(
comptime T: type,
out: Tensor(T),
image: Tensor(T),
kernel: Tensor(T),
stride: [2]usize,
pad: [2]usize,
) Error!void {
comptime requireFloat(T);
if (image.rank != 2 or kernel.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
if (stride[0] == 0 or stride[1] == 0) {
return Error.DomainError;
}
const h: usize = image.shape[0];
const w: usize = image.shape[1];
const kh: usize = kernel.shape[0];
const kw: usize = kernel.shape[1];
if (kh == 0 or kw == 0 or kh > h + 2 * pad[0] or kw > w + 2 * pad[1]) {
return Error.ShapeMismatch;
}
const oh: usize = (h + 2 * pad[0] - kh) / stride[0] + 1;
const ow: usize = (w + 2 * pad[1] - kw) / stride[1] + 1;
if (out.shape[0] != oh or out.shape[1] != ow or out.isAliased()) {
return Error.ShapeMismatch;
}
var y: usize = 0;
while (y < oh) : (y += 1) {
var x: usize = 0;
while (x < ow) : (x += 1) {
var acc: T = 0;
var i: usize = 0;
while (i < kh) : (i += 1) {
// The tap's row in the unpadded image, or out of range: signed, because the
// padding makes the top-left taps land above and left of (0, 0).
const iy: isize = @as(isize, @intCast(y * stride[0] + i)) - @as(isize, @intCast(pad[0]));
if (iy < 0 or iy >= @as(isize, @intCast(h))) {
continue;
}
var j: usize = 0;
while (j < kw) : (j += 1) {
const ix: isize = @as(isize, @intCast(x * stride[1] + j)) - @as(isize, @intCast(pad[1]));
if (ix < 0 or ix >= @as(isize, @intCast(w))) {
continue;
}
acc += (try image.at(&.{ @intCast(iy), @intCast(ix) })) * (try kernel.at(&.{ i, j }));
}
}
try out.setAt(&.{ y, x }, acc);
}
}
}
maxPool2d/// Max pooling over non-overlapping `size x size` windows. The output is `(h / size, w / size)`;
/// a trailing partial window is dropped.
///
/// Non-overlapping and square, which is what nearly every network uses. Dropping the partial
/// window rather than padding it means every output element saw exactly `size^2` inputs, so the
/// operation's meaning does not change at the edge.
pub fn maxPool2d(
comptime T: type,
out: Tensor(T),
image: Tensor(T),
size: usize,
) Error!void {
comptime requireFloat(T);
return pool2d(T, out, image, size, .max);
}
avgPool2d/// Average pooling over non-overlapping `size x size` windows.
pub fn avgPool2d(
comptime T: type,
out: Tensor(T),
image: Tensor(T),
size: usize,
) Error!void {
comptime requireFloat(T);
return pool2d(T, out, image, size, .mean);
}
The output is ((h + 2·pad − kh) / stride + 1, …). Padding is
implicit — a tap landing outside the image contributes zero and no padded copy is made
— and pad = kh/2 with stride 1 gives an output the size of the input.
conv2d slides the kernel without flipping it, which is cross-correlation.
Every framework's "conv2d" does the same: the kernel is learned, so a flipped kernel is just a
different kernel, and the unflipped form is one indexing operation cheaper. The test checks the
direction with an asymmetric kernel — [1, 0] must reproduce the left
neighbour and [0, 1] the right, which a true convolution would swap.
Rank 2, one channel. Pooling is over non-overlapping square windows and
drops a trailing partial window, so every output element saw exactly size²
inputs and the operation's meaning does not change at the edge.
var picked: [2]usize = @splat(0);
try zn.argmaxRows(f32, &picked, logits); // one class per row
const flat = try zn.argmaxAll(f32, logits); // flat index of the largest
try zn.oneHotRows(f32, hot, &.{ 2, 0 });
const loss = try zn.crossEntropyRows(f32, logits, &.{ 2, 0 });
try zn.crossEntropyRowsGrad(f32, grad, logits, &.{ 2, 0 });
const acc = try zn.accuracy(f32, logits, &.{ 2, 0 });argmaxRows/// The column of the largest element in each row of a rank-2 tensor, written into `out`.
pub fn argmaxRows(comptime T: type, out: []usize, a: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 2) {
return Error.UnsupportedShape;
}
if (out.len != a.shape[0]) {
return Error.ShapeMismatch;
}
const cols: usize = a.shape[1];
if (cols == 0) {
return Error.DomainError;
}
for (out, 0..) |*slot, row| {
var best: T = try a.at(&.{ row, 0 });
var best_at: usize = 0;
var col: usize = 1;
while (col < cols) : (col += 1) {
const x: T = a.at2(row, col);
if (x > best or (comptime isFloat(T)) and zm.isNan(best)) {
best = x;
best_at = col;
}
}
slot.* = best_at;
}
}
argmaxAll/// The flat index of the largest element. `DomainError` for an empty tensor.
///
/// Ties go to the FIRST occurrence, and NaN never wins: the comparison is `>`, so a NaN
/// candidate fails it and the running best survives. A `>=` would return the last of a run of
/// equals and would let a trailing NaN take the answer.
pub fn argmaxAll(comptime T: type, a: Tensor(T)) Error!usize {
comptime requireNumeric(T);
const n: usize = a.size();
if (n == 0) {
return Error.DomainError;
}
var walker: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
var best: T = try a.at(walker[0..a.rank]);
var best_at: usize = 0;
var i: usize = 0;
while (i < n) : (i += 1) {
const x: T = a.data[a.offsetOf(walker[0..a.rank])];
if (x > best or (comptime isFloat(T)) and zm.isNan(best)) {
best = x;
best_at = i;
}
advance(shape, walker[0..a.rank]);
}
return best_at;
}
oneHotRows/// One-hot encode class indices into the rows of `out`, which is zeroed first.
pub fn oneHotRows(comptime T: type, out: Tensor(T), indices: []const usize) Error!void {
comptime requireNumeric(T);
if (out.rank != 2 or out.isAliased()) {
return Error.UnsupportedShape;
}
if (out.shape[0] != indices.len) {
return Error.ShapeMismatch;
}
for (indices) |idx| {
if (idx >= out.shape[1]) {
return Error.OutOfRange;
}
}
out.fill(0);
for (indices, 0..) |idx, row| {
try out.setAt(&.{ row, idx }, 1);
}
}
crossEntropyRows/// Softmax cross-entropy of rank-2 `logits` against integer class `targets`, averaged over rows.
///
/// THE SOFTMAX IS NOT COMPUTED, AND THAT IS THE POINT
///
/// The definition is `-log(softmax(logits)[target])`, and computing it that way exponentiates,
/// normalises, then takes a logarithm - three chances to lose the answer. A logit of 800 gives an
/// infinite exponential and a loss of NaN; a correct probability of 1e-40 underflows to zero and
/// the logarithm is `-inf`.
///
/// The identity `-log(softmax(x)[t]) = logsumexp(x) - x[t]` avoids both: `logsumexp` is evaluated
/// by subtracting the row maximum, so the largest exponent is `e^0`, and the subtraction happens
/// in log space where the numbers are small. This form is finite for every finite input.
///
/// Takes LOGITS, not probabilities. A version taking probabilities would have to undo a softmax
/// the caller had already applied, which is where the precision is lost.
pub fn crossEntropyRows(
comptime T: type,
logits: Tensor(T),
targets: []const usize,
) Error!T {
comptime requireFloat(T);
if (logits.rank != 2) {
return Error.UnsupportedShape;
}
const rows: usize = logits.shape[0];
const cols: usize = logits.shape[1];
if (targets.len != rows) {
return Error.ShapeMismatch;
}
if (rows == 0 or cols == 0) {
return Error.DomainError;
}
for (targets) |t| {
if (t >= cols) {
return Error.OutOfRange;
}
}
var total: CompensatedSum(T) = .{};
var row: usize = 0;
while (row < rows) : (row += 1) {
var largest: T = logits.at2(row, 0);
var col: usize = 1;
while (col < cols) : (col += 1) {
largest = @max(largest, logits.at2(row, col));
}
var sum_exp: T = 0;
col = 0;
while (col < cols) : (col += 1) {
sum_exp += @exp(logits.at2(row, col) - largest);
}
const log_sum_exp: T = largest + @log(sum_exp);
total.add(log_sum_exp - logits.at2(row, targets[row]));
}
return total.value() / @as(T, @floatFromInt(rows));
}
crossEntropyRowsGrad/// The gradient of `crossEntropyRows` with respect to the logits: `softmax(logits) - onehot`,
/// divided by the number of rows.
///
/// THIS IS WHY CROSS-ENTROPY IS PAIRED WITH SOFTMAX. The two functions' derivatives compose
/// into a subtraction - no division, no `exp'`, nothing that can overflow - so the backward pass
/// of the pair is cheaper and better conditioned than either alone.
pub fn crossEntropyRowsGrad(
comptime T: type,
out: Tensor(T),
logits: Tensor(T),
targets: []const usize,
) Error!void {
comptime requireFloat(T);
if (out.rank != 2 or out.isAliased()) {
return Error.UnsupportedShape;
}
if (out.shape[0] != logits.shape[0] or out.shape[1] != logits.shape[1]) {
return Error.ShapeMismatch;
}
if (targets.len != logits.shape[0]) {
return Error.ShapeMismatch;
}
try softmaxRows(T, out, logits);
const rows: T = @floatFromInt(out.shape[0]);
for (targets, 0..) |t, row| {
if (t >= out.shape[1]) {
return Error.OutOfRange;
}
const p: T = try out.at(&.{ row, t });
try out.setAt(&.{ row, t }, p - 1);
}
var i: usize = 0;
while (i < out.size()) : (i += 1) {
out.data[out.base + i] /= rows;
}
}
accuracy/// The fraction of rows whose largest logit is at the target class.
pub fn accuracy(
comptime T: type,
logits: Tensor(T),
targets: []const usize,
) Error!T {
comptime requireFloat(T);
if (logits.rank != 2) {
return Error.UnsupportedShape;
}
if (targets.len != logits.shape[0]) {
return Error.ShapeMismatch;
}
if (targets.len == 0) {
return Error.DomainError;
}
var correct: usize = 0;
for (targets, 0..) |t, row| {
var best: T = try logits.at(&.{ row, 0 });
var best_at: usize = 0;
var col: usize = 1;
while (col < logits.shape[1]) : (col += 1) {
const x: T = logits.at2(row, col);
if (x > best or (comptime isFloat(T)) and zm.isNan(best)) {
best = x;
best_at = col;
}
}
if (best_at == t) {
correct += 1;
}
}
return @as(T, @floatFromInt(correct)) / @as(T, @floatFromInt(targets.len));
}
Class indices are []usize rather than a tensor, because that is what a caller has:
a list of labels. Targets are checked against the class count before anything is written.
argmaxAll and argmaxRows compare with >, so a tie goes
to the first occurrence and a NaN never wins — a NaN candidate fails the comparison and the
running best survives.
crossEntropyRows takes logits and never computes the softmax. The
definition is −log(softmax(x)[t]), which exponentiates, normalises, then takes a
logarithm — three chances to lose the answer. A logit of 800 gives an infinite exponential
and a loss of NaN; a correct probability of 1e-40 underflows to zero and its logarithm is
−inf.
−log(softmax(x)[t]) = logsumexp(x) − x[t] avoids both.
logsumexp subtracts the row maximum, so the largest exponent is
e0, and the subtraction happens in log space where the numbers are small.
On the row { 800, 0, −800 } this returns 0 and 800 for the two classes, both
exact, where the definition returns NaN.
crossEntropyRowsGrad is softmax(logits) − onehot, divided by the
row count. Softmax and cross-entropy are paired because their derivatives compose into a
subtraction: no division, no exponential in the backward pass, and each row of the result sums to
zero, so a step can shift probability between classes but cannot create any.
try zn.sgdStep(f32, w, w, grad, 0.1);
try zn.sgdMomentum(f32, w, w, grad, velocity, 0.1, 0.9);
try zn.adamStep(f32, w, w, grad, moment, velocity, .{}, step);sgdStep/// `out = weight - rate * grad` - one step of plain gradient descent.
///
/// NOT IN PLACE, and the caller may still pass `weight` as `out` to get in-place behaviour.
/// Returning a fresh result by default means the operation composes and can be tested against a
/// reference without the reference having been consumed; an in-place-only signature forces every
/// caller to copy first, including the tests.
///
/// Written as a loop rather than through `zip`, because `zip`'s function is a COMPTIME
/// parameter and cannot close over a runtime `rate`. A `zipScalar` seam would fix that and is
/// not worth inventing for one caller.
pub fn sgdStep(
comptime T: type,
out: Tensor(T),
weight: Tensor(T),
grad: Tensor(T),
rate: T,
) Error!void {
comptime requireFloat(T);
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const vw: Tensor(T) = try weight.broadcastTo(shape);
const vg: Tensor(T) = try grad.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const w: T = vw.data[vw.offsetOf(at)];
const g: T = vg.data[vg.offsetOf(at)];
out.data[out.offsetOf(at)] = w - rate * g;
}
}
sgdMomentum/// One SGD step with momentum. `velocity` is read and updated in place.
///
/// `velocity = momentum * velocity + grad`, then `out = weight - rate * velocity`.
///
/// THE VELOCITY CARRIES THE GRADIENT, NOT THE STEP
///
/// The other common form is `velocity = momentum * velocity + rate * grad`, which folds the
/// learning rate into the accumulator. The two differ the moment `rate` changes: with the rate
/// inside, a schedule's old rate persists in the velocity for several steps afterwards. Keeping
/// it out means a rate change takes effect immediately, which is what a schedule is for.
///
/// `momentum = 0` reduces exactly to `sgdStep`, and the test asserts it: an optimiser whose
/// degenerate case does not match the simpler one has an error in the accumulation.
pub fn sgdMomentum(
comptime T: type,
out: Tensor(T),
weight: Tensor(T),
grad: Tensor(T),
velocity: Tensor(T),
rate: T,
momentum: T,
) Error!void {
comptime requireFloat(T);
if (out.isAliased() or velocity.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const vw: Tensor(T) = try weight.broadcastTo(shape);
const vg: Tensor(T) = try grad.broadcastTo(shape);
if (velocity.size() != out.size()) {
return Error.ShapeMismatch;
}
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const v: T = velocity.data[velocity.offsetOf(at)];
const g: T = vg.data[vg.offsetOf(at)];
const updated: T = momentum * v + g;
velocity.data[velocity.offsetOf(at)] = updated;
const w: T = vw.data[vw.offsetOf(at)];
out.data[out.offsetOf(at)] = w - rate * updated;
}
}
adamStep/// One Adam step. `m` and `v` are the first and second moment buffers, read and updated in place.
/// `step` is the 1-based iteration number, which the bias correction needs.
///
/// THE BIAS CORRECTION IS WHY `step` IS A PARAMETER
///
/// `m` and `v` start at zero, so the first few updates are biased toward zero by exactly the
/// factor `1 - beta^step`. Dividing it out makes the FIRST step the same size as any other; an
/// implementation without the correction takes an opening step smaller by a factor of about
/// `1 - beta2`, which for the usual 0.999 means a thousandth of the intended distance.
///
/// THE STEP SIZE IS INDEPENDENT OF THE GRADIENT'S SCALE. With `m` and `v` both proportional to
/// the gradient, the ratio `m / sqrt(v)` is not, so multiplying every gradient by a thousand
/// leaves the update almost unchanged. That property is what the test checks, because it is the
/// reason to use Adam at all.
pub fn adamStep(
comptime T: type,
out: Tensor(T),
weight: Tensor(T),
grad: Tensor(T),
moment: Tensor(T),
velocity: Tensor(T),
hyper: Adam,
step: usize,
) Error!void {
comptime requireFloat(T);
if (out.isAliased() or moment.isAliased() or velocity.isAliased()) {
return Error.UnsupportedShape;
}
if (step == 0) {
return Error.DomainError;
}
const shape: []const usize = out.shape[0..out.rank];
const vw: Tensor(T) = try weight.broadcastTo(shape);
const vg: Tensor(T) = try grad.broadcastTo(shape);
if (moment.size() != out.size() or velocity.size() != out.size()) {
return Error.ShapeMismatch;
}
const beta1: T = @floatCast(hyper.beta1);
const beta2: T = @floatCast(hyper.beta2);
const eps: T = @floatCast(hyper.epsilon);
const rate: T = @floatCast(hyper.rate);
// EXPONENTIATION BY SQUARING, NOT REPEATED MULTIPLICATION
//
// The first version of this multiplied `step` times to avoid "a rounding `pow` would
// introduce". That traded an unmeasured rounding for **O(step) work on every call** - 50
// million multiplications over a 10 000-step run, 5 billion over 100 000, and growing
// quadratically with the length of training. znum uses `pow` and is right to.
//
// Squaring is O(log step) - 17 multiplications at 100 000 steps rather than 100 000 - and
// stays exact in the sense that matters: it is the same sequence of floating-point
// multiplications, just associated differently.
const correct1: T = 1 - powInt(T, beta1, step);
const correct2: T = 1 - powInt(T, beta2, step);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const g: T = vg.data[try vg.flatIndex(at)];
const m_index: usize = try moment.flatIndex(at);
const v_index: usize = try velocity.flatIndex(at);
const m_new: T = beta1 * moment.data[m_index] + (1 - beta1) * g;
const v_new: T = beta2 * velocity.data[v_index] + (1 - beta2) * g * g;
moment.data[m_index] = m_new;
velocity.data[v_index] = v_new;
const m_hat: T = m_new / correct1;
const v_hat: T = v_new / correct2;
const w: T = vw.data[vw.offsetOf(at)];
out.data[out.offsetOf(at)] =
w - rate * m_hat / (@sqrt(v_hat) + eps);
}
}
Each writes into a destination that may be the weight itself. The state buffers —
velocity for momentum, moment and velocity for Adam —
are read and updated in place, are the caller's to allocate, and must start at zero.
sgdMomentum accumulates the gradient, not the step:
velocity = momentum * velocity + grad
out = weight - rate * velocityThe alternative folds rate into the accumulator. The two differ as soon as the rate
changes: with the rate inside, an old rate persists in the velocity for several steps afterwards.
With momentum = 0 this reduces exactly to sgdStep.
adamStep takes its hyperparameters as a struct with the paper's defaults, and a
1-based step number:
pub const Adam = struct {
rate: f64 = 0.001,
beta1: f64 = 0.9,
beta2: f64 = 0.999,
epsilon: f64 = 1.0e-8,
};step is a parameter because of the bias correction. moment and
velocity start at zero, so early updates are biased toward zero by exactly
1 − betastep; dividing that out makes the first step the same size as
any other. Without it the opening step is smaller by roughly 1 − beta2, which at
the default 0.999 is a thousandth of the intended distance. A step of 0 is
DomainError.
The reason to use Adam is that the step size does not depend on the gradient's scale. Both moments are proportional to the gradient, so their ratio is not: multiplying every gradient by a thousand leaves the update essentially unchanged.
const before = try zn.clipByNorm(f32, grad, grad, 1.0);clipByNorm/// Scale `a` into `out` so its Euclidean norm is at most `limit`. Returns the norm BEFORE
/// clipping, which is the number worth logging.
///
/// THE WHOLE TENSOR IS SCALED BY ONE FACTOR, not clipped elementwise. Clipping each element
/// separately changes the gradient's DIRECTION, which is the one property a gradient has that
/// must survive: a step in a different direction is not a smaller version of the step that was
/// wanted.
pub fn clipByNorm(comptime T: type, out: Tensor(T), a: Tensor(T), limit: T) Error!T {
comptime requireFloat(T);
if (limit <= 0) {
return Error.DomainError;
}
const total: T = try norm(T, a);
if (total <= limit) {
try materialise(T, out, a);
return total;
}
try scale(T, out, a, limit / total);
return total;
}
Scales the whole tensor by a single factor so its Euclidean norm is at most the limit, and returns the norm before clipping — the number worth logging. Below the limit the data is copied unchanged.
One factor for the whole tensor, rather than clipping each element, because elementwise clipping changes the gradient's direction. A step in a different direction is not a smaller version of the step that was wanted.
Every operation a single dense layer needs is present. With x the input,
w the weights and bias a row vector:
// forward: y = relu(x @ w + bias)
try zn.matmul(f32, pre, x, w);
try zn.add(f32, pre, pre, try bias.broadcastTo(pre.shape[0..pre.rank]));
try zn.relu(f32, y, pre);
const loss = try zn.mseLoss(f32, y, target);
// backward
try zn.sub(f32, dy, y, target); // d(loss)/dy, up to a constant
try zn.reluGrad(f32, dpre, pre, dy); // through the activation
try zn.matmul(f32, dw, try x.transpose(0, 1), dpre); // dW = xT @ dpre
try zn.sumAxis(f32, dbias, dpre, 0); // bias gradient
try zn.matmul(f32, dx, dpre, try w.transpose(0, 1)); // dx = dpre @ wT
// update
try zn.sgdStep(f32, w, w, dw, rate);matmul/// `out = a @ b` - the matrix product, for rank-2 tensors only.
///
/// THE ONE OPERATION THAT IS NOT ELEMENTWISE
///
/// `(m, k) @ (k, n) -> (m, n)`. Every shape is checked; a mismatch is `ShapeMismatch` before any
/// element is written, so a partially-filled output is not a state this can produce.
///
/// The inner sum is compensated, for the reason in `sumAll`: a dot product over a long
/// contraction axis is exactly the accumulation that loses its low bits, and in a neural network
/// that axis is the layer width. Getting this wrong does not crash - it makes gradients slightly
/// wrong, everywhere, in a way that reads as "training is a bit unstable".
///
/// This is the naive triple loop, deliberately. It is the SPECIFICATION every faster
/// implementation is checked against, and it stays in the file when a blocked or tiled version
/// arrives. Strided and stretched inputs are handled by going through `at`, so a transposed view
/// multiplies correctly without being copied first.
pub fn matmul(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
comptime requireFloat(T);
if (a.rank != 2 or b.rank != 2 or out.rank != 2) {
return Error.UnsupportedShape;
}
const m: usize = a.shape[0];
const k: usize = a.shape[1];
const n: usize = b.shape[1];
if (b.shape[0] != k or out.shape[0] != m or out.shape[1] != n) {
return Error.ShapeMismatch;
}
if (out.isAliased()) {
return Error.UnsupportedShape;
}
var row: usize = 0;
while (row < m) : (row += 1) {
var col: usize = 0;
while (col < n) : (col += 1) {
var total: CompensatedSum(T) = .{};
var i: usize = 0;
while (i < k) : (i += 1) {
const product: T = a.at2(row, i) *
b.at2(i, col);
total.add(product);
}
out.setAt2(row, col, total.value());
}
}
}
add/// `out = a + b`, broadcasting both to `out`'s shape.
pub fn add(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return x + y;
}
}.apply);
}
relu/// `out = max(a, 0)`, elementwise.
pub fn relu(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
return map(T, out, a, struct {
fn apply(x: T) T {
return if (x > 0) x else 0;
}
}.apply);
}
mseLoss/// Mean squared error between two tensors of the same shape.
///
/// Returns a scalar rather than writing into a tensor, because a loss IS a scalar and handing
/// back a rank-0 tensor would make every caller unwrap it. `DomainError` for empty input, since
/// the mean of nothing is a question for the caller.
pub fn mseLoss(comptime T: type, prediction: Tensor(T), target: Tensor(T)) Error!T {
comptime requireFloat(T);
if (prediction.rank != target.rank) {
return Error.ShapeMismatch;
}
var axis: usize = 0;
while (axis < prediction.rank) : (axis += 1) {
if (prediction.shape[axis] != target.shape[axis]) {
return Error.ShapeMismatch;
}
}
const n: usize = prediction.size();
if (n == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var walk: Walk = .over(prediction.shape[0..prediction.rank]);
while (walk.next()) |at| {
const delta: T = (try prediction.at(at)) - (try target.at(at));
const squared: T = delta * delta;
total.add(squared);
}
const divisor: T = @floatFromInt(n);
return total.value() / divisor;
}
sub/// `out = a - b`, broadcasting both to `out`'s shape.
pub fn sub(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void {
return zip(T, out, a, b, struct {
fn apply(x: T, y: T) T {
return x - y;
}
}.apply);
}
reluGrad/// `out = if (x > 0) grad else 0` - the gradient of `relu` with respect to its input.
///
/// It takes the FORWARD INPUT, not the forward output. Both work for `relu` because the output
/// is positive exactly where the input was, but only the input generalises: `leakyRelu` and
/// `elu` have outputs that do not determine their own gradient. Taking the input here means the
/// signature does not have to change when they arrive.
pub fn reluGrad(
comptime T: type,
out: Tensor(T),
x: Tensor(T),
grad: Tensor(T),
) Error!void {
comptime requireFloat(T);
return zip(T, out, x, grad, struct {
fn apply(xv: T, gv: T) T {
return if (xv > 0) gv else 0;
}
}.apply);
}
sumAxis/// Sum `a` along `axis` into `out`, which must have `a`'s shape with that axis removed.
///
/// The output's rank is one less, not one with a size-1 axis. Keeping the axis would be the
/// `keepdims` option, and options that change the RESULT SHAPE are better as separate calls than
/// as a flag a reader has to look up - a caller who wants the axis back has `reshape`.
pub fn sumAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void {
comptime requireNumeric(T);
if (axis >= a.rank) {
return Error.OutOfRange;
}
if (out.rank + 1 != a.rank or out.isAliased()) {
return Error.ShapeMismatch;
}
// The surviving axes must match, in order.
var src: usize = 0;
var dst: usize = 0;
while (src < a.rank) : (src += 1) {
if (src == axis) {
continue;
}
if (a.shape[src] != out.shape[dst]) {
return Error.ShapeMismatch;
}
dst += 1;
}
out.fill(0);
var walker: [max_rank]usize = @splat(0);
var reduced: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
var remaining: usize = a.size();
while (remaining > 0) : (remaining -= 1) {
// Project the source position onto the output by dropping the reduced axis.
var k: usize = 0;
var i: usize = 0;
while (i < a.rank) : (i += 1) {
if (i == axis) {
continue;
}
reduced[k] = walker[i];
k += 1;
}
const at: []const usize = reduced[0..out.rank];
const slot: usize = out.offsetOf(at);
out.data[slot] += a.data[a.offsetOf(walker[0..a.rank])];
advance(shape, walker[0..a.rank]);
}
}
sgdStep/// `out = weight - rate * grad` - one step of plain gradient descent.
///
/// NOT IN PLACE, and the caller may still pass `weight` as `out` to get in-place behaviour.
/// Returning a fresh result by default means the operation composes and can be tested against a
/// reference without the reference having been consumed; an in-place-only signature forces every
/// caller to copy first, including the tests.
///
/// Written as a loop rather than through `zip`, because `zip`'s function is a COMPTIME
/// parameter and cannot close over a runtime `rate`. A `zipScalar` seam would fix that and is
/// not worth inventing for one caller.
pub fn sgdStep(
comptime T: type,
out: Tensor(T),
weight: Tensor(T),
grad: Tensor(T),
rate: T,
) Error!void {
comptime requireFloat(T);
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const vw: Tensor(T) = try weight.broadcastTo(shape);
const vg: Tensor(T) = try grad.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const w: T = vw.data[vw.offsetOf(at)];
const g: T = vg.data[vg.offsetOf(at)];
out.data[out.offsetOf(at)] = w - rate * g;
}
}
There is no autograd; the chain above is written by hand.
try zn.adamWStep(f64, out, w, g, m, v, hyper, decay, step);
try zn.rmspropStep(f64, out, w, g, mean_square, rate, rho, eps);
try zn.adagradStep(f64, out, w, g, acc, rate, eps);
try zn.clipByValue(f64, out, g, limit);
const lr = zn.cosineLearningRate(f64, step, total, base, lowest);
const wu = zn.warmupLearningRate(f64, step, warmup, base);adamWStep/// One AdamW step: Adam, plus weight decay applied DIRECTLY TO THE WEIGHT.
///
/// WHAT MAKES THIS AdamW AND NOT ADAM WITH L2
///
/// The old way of shrinking weights is to add `decay * weight` to the gradient. With plain SGD
/// the two are the same thing. With Adam they are not: the gradient goes through the moment
/// estimates, so the decay is divided by `sqrt(v_hat)` along with everything else - **a
/// parameter with a large historical gradient gets decayed LESS than one with a small one**,
/// which is the opposite of what a regulariser is for. The decay ends up depending on the
/// gradient history of the weight it is supposed to shrink uniformly.
///
/// Decoupling means subtracting `rate * decay * weight` from the weight after the Adam update,
/// where no moment can touch it. znum does this and documents why; this follows it.
///
/// `decay = 0` reduces exactly to `adamStep`, which the test asserts - an implementation whose
/// degenerate case disagrees with the simpler one has the coupling in the wrong place.
pub fn adamWStep(
comptime T: type,
out: Tensor(T),
weight: Tensor(T),
grad: Tensor(T),
moment: Tensor(T),
velocity: Tensor(T),
hyper: Adam,
decay: T,
step: usize,
) Error!void {
comptime requireFloat(T);
if (decay < 0) {
return Error.DomainError;
}
try adamStep(T, out, weight, grad, moment, velocity, hyper, step);
if (decay == 0) {
return;
}
const rate: T = @floatCast(hyper.rate);
const shape: []const usize = out.shape[0..out.rank];
const weight_view: Tensor(T) = try weight.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const w_value: T = try weight_view.at(at);
const slot: usize = try out.flatIndex(at);
out.data[slot] -= rate * decay * w_value;
}
}
rmspropStep/// One RMSProp step: divide the gradient by the root of its own running mean square.
///
/// `mean_square = rho * mean_square + (1 - rho) * grad^2`, then
/// `out = weight - rate * grad / (sqrt(mean_square) + epsilon)`.
///
/// Adam without the first moment, and without bias correction. The missing correction is why
/// RMSProp's opening steps are small where Adam's are not - `mean_square` starts at zero, so the
/// first division is by something too small and the step is too large, then it settles. Stated
/// rather than corrected, because correcting it would make this a different algorithm wearing
/// RMSProp's name.
pub fn rmspropStep(
comptime T: type,
out: Tensor(T),
weight: Tensor(T),
grad: Tensor(T),
mean_square: Tensor(T),
rate: T,
rho: T,
epsilon: T,
) Error!void {
comptime requireFloat(T);
if (rho < 0 or rho >= 1 or epsilon <= 0) {
return Error.DomainError;
}
if (out.isAliased() or mean_square.isAliased()) {
return Error.UnsupportedShape;
}
if (mean_square.size() != out.size()) {
return Error.ShapeMismatch;
}
const shape: []const usize = out.shape[0..out.rank];
const weight_view: Tensor(T) = try weight.broadcastTo(shape);
const grad_view: Tensor(T) = try grad.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const g_value: T = try grad_view.at(at);
const index: usize = try mean_square.flatIndex(at);
const updated: T = rho * mean_square.data[index] + (1 - rho) * g_value * g_value;
mean_square.data[index] = updated;
const w_value: T = try weight_view.at(at);
out.data[try out.flatIndex(at)] = w_value - rate * g_value / (@sqrt(updated) + epsilon);
}
}
adagradStep/// One Adagrad step: divide by the root of the ACCUMULATED sum of squared gradients.
///
/// The accumulator only grows, so the effective learning rate only shrinks - monotonically, to
/// zero. That is the algorithm's strength on sparse problems, where a rarely-seen feature keeps
/// a large step long after common ones have decayed, and its weakness on dense ones, where
/// learning stops before the model has converged. RMSProp's `rho` exists precisely to forget.
pub fn adagradStep(
comptime T: type,
out: Tensor(T),
weight: Tensor(T),
grad: Tensor(T),
accumulator: Tensor(T),
rate: T,
epsilon: T,
) Error!void {
comptime requireFloat(T);
if (epsilon <= 0) {
return Error.DomainError;
}
if (out.isAliased() or accumulator.isAliased()) {
return Error.UnsupportedShape;
}
if (accumulator.size() != out.size()) {
return Error.ShapeMismatch;
}
const shape: []const usize = out.shape[0..out.rank];
const weight_view: Tensor(T) = try weight.broadcastTo(shape);
const grad_view: Tensor(T) = try grad.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const g_value: T = try grad_view.at(at);
const index: usize = try accumulator.flatIndex(at);
const total: T = accumulator.data[index] + g_value * g_value;
accumulator.data[index] = total;
const w_value: T = try weight_view.at(at);
out.data[try out.flatIndex(at)] = w_value - rate * g_value / (@sqrt(total) + epsilon);
}
}
clipByValue/// Clip every element of `a` into `[-limit, limit]`, into `out`.
///
/// ELEMENTWISE, WHICH IS THE OPPOSITE CHOICE FROM `clipByNorm` AND DELIBERATELY SO. This changes
/// the gradient's direction, because it shortens some components and not others. That is
/// sometimes what you want - a single exploding coordinate is contained without touching the
/// rest - and it is never what you want when the direction matters. `clipByNorm` is the one that
/// preserves direction; having both under names that say which is which is the point.
pub fn clipByValue(comptime T: type, out: Tensor(T), a: Tensor(T), limit: T) Error!void {
comptime requireNumeric(T);
if (limit <= 0) {
return Error.DomainError;
}
return clamp(T, out, a, -limit, limit);
}
cosineLearningRate/// The learning rate at `step` under a cosine schedule from `base` down to `lowest`.
///
/// Follows znum exactly: progress is `step / total`, and the rate is
/// `lowest + (base - lowest) * (1 + cos(pi * progress)) / 2`. At step 0 that is `base`, at
/// `total` it is `lowest`, and the descent is steepest in the middle - which is the shape the schedule is
/// chosen for. At or past `total` it returns `lowest` rather than continuing the cosine into its
/// rising half, which would take the rate back up.
pub fn cosineLearningRate(
comptime T: type,
step: usize,
total: usize,
base: T,
lowest: T,
) T {
comptime requireFloat(T);
// `lowest`, not `floor`: this file exports a tensor `floor`, and a parameter of that name
// would shadow it. The same shadowing the rename found four times elsewhere.
if (total == 0 or step >= total) {
return lowest;
}
const progress: T = @as(T, @floatFromInt(step)) / @as(T, @floatFromInt(total));
const half_turn: T = @floatCast(3.141592653589793);
return lowest + (base - lowest) * (1 + scalarCos(half_turn * progress)) / 2;
}
warmupLearningRate/// A linear ramp from zero to `base` over the first `warmup` steps, then `base`.
///
/// WHY A WARMUP EXISTS AT ALL. Adam's second moment starts at zero and its first few estimates
/// are made from one or two gradients, so the early updates are both large and badly informed.
/// Ramping the rate in means the moments have seen enough data to be worth dividing by before
/// the steps get big. It is the difference between a transformer that trains and one that
/// diverges in the first hundred steps.
///
/// Step 0 gives exactly 0, and step `warmup` gives exactly `base`. `warmup = 0` is the constant
/// schedule rather than a division by zero.
pub fn warmupLearningRate(comptime T: type, step: usize, warmup: usize, base: T) T {
comptime requireFloat(T);
if (warmup == 0 or step >= warmup) {
return base;
}
return base * @as(T, @floatFromInt(step)) / @as(T, @floatFromInt(warmup));
}
The old way of shrinking weights adds decay · weight to the gradient. With
plain SGD that is the same thing as decaying the weight. With Adam it is not: the
gradient goes through the moment estimates, so the decay is divided by
sqrt(v̂) along with everything else — a parameter with a large historical
gradient gets decayed less than one with a small one, which is the opposite of what a
regulariser is for.
So there is no claim that the weights got smaller, which any wrong version would also
pass. It checks that the difference from plain Adam is exactly
rate · decay · weight — a quantity depending on the
weight alone. Had the decay gone through the moments, the difference would depend on the gradient
too, and that equality would fail on the elements whose gradients differ most.
rmspropStep and adagradStep differ in one thing: whether they forget.
Fed the same gradient forty times, Adagrad's step shrinks from 1.00e-1 to 1.58e-2 and never stops
shrinking, while RMSProp's settles at 1.01e-1. The measurement also shows what RMSProp's missing
bias correction costs: its first step is 3.16e-1, three times its settled one,
because mean_square starts at zero. Adam corrects for that; RMSProp does not, and
that is a property of the algorithm rather than of this implementation.
clipByValue is elementwise and therefore changes the gradient's
direction; clipByNorm scales the whole tensor and preserves it. Both are
wanted at different times and the distinction directly — every ratio to
the original is equal after norm clipping, and visibly unequal after value clipping.
try zn.confusionMatrix(f64, out, predicted, truth); // true DOWN, predicted ACROSS
const s = try zn.classScore(f64, confusion, k);
const macro = try zn.macroF1(f64, confusion);
const r2 = try zn.r2Score(f64, prediction, truth); // can be NEGATIVE
const cos = try zn.cosineSimilarity(f64, a, b);
const hinge = try zn.hingeLoss(f64, scores, labels);confusionMatrix/// The confusion matrix of predicted against true classes: `out[t][p]` counts the samples whose
/// true class is `t` and predicted class is `p`.
///
/// TRUE DOWN, PREDICTED ACROSS, which is scikit-learn's orientation and znum's. The transpose is
/// equally common elsewhere and reading a confusion matrix the wrong way round swaps precision
/// with recall, so the convention is stated rather than left to the reader's habit.
pub fn confusionMatrix(
comptime T: type,
out: Tensor(T),
predicted: []const usize,
truth: []const usize,
) Error!void {
comptime requireFloat(T);
if (out.rank != 2 or out.shape[0] != out.shape[1] or out.isAliased()) {
return Error.UnsupportedShape;
}
if (predicted.len != truth.len) {
return Error.ShapeMismatch;
}
const classes: usize = out.shape[0];
for (predicted, truth) |p, t| {
if (p >= classes or t >= classes) {
return Error.OutOfRange;
}
}
out.fill(0);
for (predicted, truth) |p, t| {
try out.setAt(&.{ t, p }, (try out.at(&.{ t, p })) + 1);
}
}
classScore/// The scores for class `k` of a confusion matrix laid out true-down, predicted-across.
pub fn classScore(comptime T: type, confusion: Tensor(T), k: usize) Error!ClassScore {
comptime requireFloat(T);
if (confusion.rank != 2 or confusion.shape[0] != confusion.shape[1]) {
return Error.UnsupportedShape;
}
const classes: usize = confusion.shape[0];
if (k >= classes) {
return Error.OutOfRange;
}
const hit: f64 = @floatCast(try confusion.at(&.{ k, k }));
var predicted_total: f64 = 0;
var true_total: f64 = 0;
for (0..classes) |i| {
predicted_total += @floatCast(try confusion.at(&.{ i, k })); // column: predicted as k
true_total += @floatCast(try confusion.at(&.{ k, i })); // row: truly k
}
const precision: f64 = if (predicted_total > 0) hit / predicted_total else 0;
const recall: f64 = if (true_total > 0) hit / true_total else 0;
const denominator: f64 = precision + recall;
return .{
.precision = precision,
.recall = recall,
.f1 = if (denominator > 0) 2 * precision * recall / denominator else 0,
.precision_undefined = predicted_total == 0,
.recall_undefined = true_total == 0,
};
}
macroF1/// The unweighted mean of every class's F1 - the macro average.
///
/// Unweighted, so a class with three samples counts as much as one with three thousand. That is
/// the point of a macro average and also its trap: on an imbalanced problem it is dominated by
/// the rare classes, which is either exactly what you want or exactly what misleads you. The
/// weighted alternative is `accuracy`, which is already here.
pub fn macroF1(comptime T: type, confusion: Tensor(T)) Error!f64 {
comptime requireFloat(T);
const classes: usize = confusion.shape[0];
if (classes == 0) {
return Error.DomainError;
}
var total: f64 = 0;
for (0..classes) |k| {
total += (try classScore(T, confusion, k)).f1;
}
const class_count: f64 = @floatFromInt(classes);
return total / class_count;
}
r2Score/// The coefficient of determination: `1 - residual_sum_of_squares / total_sum_of_squares`.
///
/// IT CAN BE NEGATIVE, AND THAT IS INFORMATION
///
/// The name says "R squared" and squares are not negative, which misleads almost everyone who
/// meets it. `1 - ss_res/ss_tot` is negative whenever the model does worse than predicting the
/// mean of the truth - which happens, and is exactly what you want to be told. Clamping it at
/// zero would turn "worse than a constant" into "explains nothing", and those are different.
///
/// `DomainError` when the truth is constant: `ss_tot` is then zero, and the question "what
/// fraction of the variation is explained" has no answer when there is no variation.
pub fn r2Score(comptime T: type, prediction: Tensor(T), truth: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = truth.size();
if (n != prediction.size()) {
return Error.ShapeMismatch;
}
if (n == 0) {
return Error.DomainError;
}
const mean: T = try meanAll(T, truth);
var residual: CompensatedSum(T) = .{};
var total: CompensatedSum(T) = .{};
var pair: Walk2 = .over(prediction.shape[0..prediction.rank], truth.shape[0..truth.rank]);
while (pair.next()) |at| {
const predicted: T = try prediction.at(at.first);
const observed: T = try truth.at(at.second);
residual.add((observed - predicted) * (observed - predicted));
total.add((observed - mean) * (observed - mean));
}
const denominator: T = total.value();
if (denominator == 0) {
return Error.DomainError;
}
return 1 - residual.value() / denominator;
}
cosineSimilarity/// The cosine of the angle between two tensors read as vectors.
///
/// `DomainError` for a zero vector, which has no direction to compare. Returning 0 would mean
/// "perpendicular", and a zero vector is not perpendicular to anything - it is undefined, and
/// silently calling it orthogonal is how a degenerate embedding becomes a plausible-looking
/// similarity score.
pub fn cosineSimilarity(comptime T: type, a: Tensor(T), b: Tensor(T)) Error!T {
comptime requireFloat(T);
if (a.size() != b.size()) {
return Error.ShapeMismatch;
}
const norm_a: T = try norm(T, a);
const norm_b: T = try norm(T, b);
if (norm_a == 0 or norm_b == 0) {
return Error.DomainError;
}
return (try dotAll(T, a, b)) / (norm_a * norm_b);
}
hingeLoss/// Hinge loss for binary labels in `{-1, +1}`: `mean(max(0, 1 - label * score))`.
///
/// The margin is 1, not 0: a sample classified correctly but only just still contributes, which
/// is what pushes the boundary away from the data rather than merely onto the right side of it.
/// That margin is the whole difference between a hinge loss and a step function.
pub fn hingeLoss(comptime T: type, scores: Tensor(T), labels: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = scores.size();
if (n != labels.size()) {
return Error.ShapeMismatch;
}
if (n == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var pair: Walk2 = .over(scores.shape[0..scores.rank], labels.shape[0..labels.rank]);
while (pair.next()) |at| {
const score: T = try scores.at(at.first);
const label: T = try labels.at(at.second);
total.add(@max(0, 1 - label * score));
}
return total.value() / @as(T, @floatFromInt(n));
}
Precision is "of everything called class k, how much was". If the model never predicted k at all, that denominator is zero and the quantity is genuinely undefined — there is no fraction of predictions that were right when there were no predictions. znum returns 0 and so does this, because a macro average has to average something. But 0 is a lie in a specific way: it says every prediction of this class was wrong, when the truth is that there were none.
So ClassScore carries precision_undefined and
recall_undefined alongside the numbers. The average uses the plain zeros; a caller
who cares can tell the two kinds of zero apart. The test is built around a class that occurs once
in the truth and is never predicted, so recall is a real zero — the model missed
it — while precision's zero describes nothing.
r2Score can be negative, and that is information. The name says
"R squared" and squares are not negative, which misleads nearly everyone who meets it.
1 − ss_res/ss_tot goes below zero whenever the model does worse than predicting
the mean, which happens and is exactly what you want to be told. Clamping at zero would turn
"worse than a constant" into "explains nothing", and those are different. A constant truth is
DomainError: there is no variation to explain.
cosineSimilarity refuses a zero vector rather than returning 0. Zero would mean
"perpendicular", and a zero vector is not perpendicular to anything — silently calling it
orthogonal is how a degenerate embedding becomes a plausible-looking similarity score.
const Net = zn.Chain(f64, struct {
first: zn.Dense(f64),
second: zn.Dense(f64),
});
const run = try net.attach(&graph, x);
var opt = try zn.Optimizer(f64).init(gpa, &graph, &run.parameters, .{ .adam = ... });Chain/// Layers run in order, with their parameters collected for you.
///
/// COLLECTING THE PARAMETERS IS THE POINT; RUNNING THEM IN ORDER IS THE EASY HALF
///
/// Writing the forward pass by hand is fine - three calls, three named intermediates, and the
/// names help when a shape is wrong. What is NOT fine is the line after it:
///
/// &.{ a1.weight, a1.bias, a2.weight, a2.bias }
///
/// That list is written by hand today, and **a parameter left out of it silently never trains**.
/// The loss still falls, because the other parameters compensate; the network is simply worse
/// than it should be, and nothing anywhere says so. It is the quietest bug in the file.
///
/// A `Chain` holds its layers in a tuple and hands back every parameter it attached, in order,
/// so the list cannot disagree with the network. znum's version iterates a struct's field types
/// and carries a fused-activation tag table; this one carries a tuple and a count, because the
/// tape already fuses nothing and the graph already composes.
pub fn Chain(comptime T: type, comptime Layers: type) type {
return struct {
layers: Layers,
const Self = @This();
const info = @typeInfo(Layers).@"struct";
/// Every parameter this chain attached, in the order the layers were listed.
///
/// A fixed array rather than a slice: the size is known at comptime from the layer list,
/// so there is nothing to allocate and nothing to free.
pub const Parameters = [parameterCount()]Var;
/// What `attach` hands back: the last layer's output, and every parameter on the way.
pub const Run = struct {
out: Var,
parameters: Parameters,
};
/// How many `Var`s `attach` will produce, counted at comptime from the layer types.
///
/// Every layer declares `pub const Attached`, whose fields are `out` plus one per
/// parameter - so the count is the field count less one, and it comes from the layer
/// ITSELF rather than from a table here that could fall out of step with it.
fn parameterCount() usize {
var total: usize = 0;
for (info.field_types) |Layer| {
total += @typeInfo(Layer.Attached).@"struct".field_names.len - 1;
}
return total;
}
/// Run every layer in order and collect what they attached.
pub fn attach(self: Self, graph: *Graph(T), input: Var) Error!Run {
var current: Var = input;
var collected: Parameters = undefined;
comptime var next: usize = 0;
inline for (info.field_names) |layer_name| {
const attached = try @field(self.layers, layer_name).attach(graph, current);
current = attached.out;
inline for (@typeInfo(@TypeOf(attached)).@"struct".field_names) |slot| {
if (comptime !std.mem.eql(u8, slot, "out")) {
collected[next] = @field(attached, slot);
next += 1;
}
}
}
return .{ .out = current, .parameters = collected };
}
};
}
Densepub fn Dense(comptime T: type) type {
comptime requireFloat(T);
return struct {
const Self = @This();
weight: Tensor(T),
bias: Tensor(T),
/// Allocate and initialise. `He` for a layer feeding a rectifier, `Xavier` otherwise.
pub fn init(
gpa: Allocator,
rng: Rng,
fan_in: usize,
fan_out: usize,
scheme: InitScheme,
) Error!Self {
const weight: Tensor(T) = try Tensor(T).alloc(gpa, &.{ fan_in, fan_out });
const bias: Tensor(T) = try Tensor(T).alloc(gpa, &.{ 1, fan_out });
switch (scheme) {
.xavier => try initXavier(T, weight, rng, fan_in, fan_out),
.he => try initHe(T, weight, rng, fan_in),
}
// Biases start at zero. A random bias adds nothing a random weight does not already
// provide, and a non-zero one shifts every unit's activation before any data is seen.
bias.fill(0);
return .{ .weight = weight, .bias = bias };
}
/// What `attach` hands back: the output, and handles for the two parameters.
///
/// RETURNING ONLY THE OUTPUT IS NOT ENOUGH, which a first attempt at a training loop
/// showed immediately: the loop needs the parameter handles to read their gradients, and
/// with the layer holding no `Var` there is nowhere else to get them. Returning all three
/// keeps the layer graph-agnostic AND gives the caller what a step requires.
pub const Attached = struct {
out: Var,
weight: Var,
bias: Var,
};
/// Record `input @ weight + bias` onto `graph`.
pub fn attach(self: Self, graph: *Graph(T), input: Var) Error!Attached {
const w: Var = try graph.parameter(self.weight);
const b: Var = try graph.parameter(self.bias);
return .{
.out = try graph.add(try graph.matmul(input, w), b),
.weight = w,
.bias = b,
};
}
};
}
Optimizer/// An optimiser bound to a set of a graph's parameters, owning whatever state the algorithm
/// needs per parameter.
///
/// THE STATE LIVES HERE, NOT WITH THE CALLER
///
/// `sgdMomentum` needs a velocity per parameter and `adamStep` needs two moments; without this,
/// a caller training four parameters with Adam allocates and threads eight extra tensors by
/// hand, and a mistake - one moment buffer shared between two parameters - trains silently
/// wrong. Binding the state to the parameter list at construction makes that mistake
/// unrepresentable.
///
/// ONE TYPE, THREE ALGORITHMS. The choice is a field, not a type parameter, so a training loop
/// can be written once and the algorithm changed at one call site. The state buffers for an
/// algorithm that does not use them are never allocated.
///
/// `step` reads gradients through the graph and writes updates into the graph's parameter
/// tensors, so the next `recompute` sees them with no further bookkeeping.
pub fn Optimizer(comptime T: type) type {
comptime requireFloat(T);
return struct {
const Self = @This();
pub const Kind = union(enum) {
sgd: struct { rate: T },
momentum: struct { rate: T, momentum: T },
adam: Adam,
};
kind: Kind,
params: []const Var,
/// Per parameter: one buffer for momentum, two for Adam, none for SGD.
state: []Tensor(T),
/// Adam's 1-based step counter.
steps: usize,
pub fn init(
gpa: Allocator,
graph: *const Graph(T),
params: []const Var,
kind: Kind,
) Error!Self {
const per: usize = switch (kind) {
.sgd => 0,
.momentum => 1,
.adam => 2,
};
// THE PARAMETER LIST IS COPIED, NOT RETAINED
//
// The first version stored the caller's slice. Every test passed, because a test
// writes `&.{ a, b, c }` in the same frame that calls `step`. The GPU training example
// built its optimiser in `init` and stepped it from `update`, and the array literal
// was gone - an out-of-bounds read inside `valueOf`, on the first frame. A type that
// outlives the call that made it owns what it needs.
const owned: []Var = try gpa.alloc(Var, params.len);
@memcpy(owned, params);
const state: []Tensor(T) = try gpa.alloc(Tensor(T), params.len * per);
for (params, 0..) |v, i| {
const shape: []const usize = graph.valueOf(v).shape[0..graph.valueOf(v).rank];
var k: usize = 0;
while (k < per) : (k += 1) {
state[i * per + k] = try Tensor(T).alloc(gpa, shape);
state[i * per + k].fill(0);
}
}
return .{ .kind = kind, .params = owned, .state = state, .steps = 0 };
}
/// Apply one update to every bound parameter, from the gradients `backward` left.
pub fn step(self: *Self, graph: *const Graph(T)) Error!void {
self.steps += 1;
for (self.params, 0..) |v, i| {
const weight: Tensor(T) = graph.valueOf(v);
const grad: Tensor(T) = try graph.gradOf(v);
switch (self.kind) {
.sgd => |h| try sgdStep(T, weight, weight, grad, h.rate),
.momentum => |h| try sgdMomentum(
T,
weight,
weight,
grad,
self.state[i],
h.rate,
h.momentum,
),
.adam => |h| try adamStep(
T,
weight,
weight,
grad,
self.state[i * 2],
self.state[i * 2 + 1],
h,
self.steps,
),
}
}
}
};
}
Writing the forward pass by hand is fine — three calls, three named intermediates, and the names help when a shape is wrong. The line after it is not:
&.{ a1.weight, a1.bias, a2.weight, a2.bias }That list is written by hand, and a parameter missing from it silently never
trains. The loss still falls — the other parameters compensate — so the network is simply
worse than it should be and nothing anywhere says so. It is the quietest bug in this library, and
it is the reason Chain exists. Running the layers in order is the easy half.
The count comes from each layer's own pub const Attached, whose fields are
out plus one per parameter. So adding a layer changes the count without anyone editing
a table, and a layer that gained a parameter would be counted correctly the day it did.
Parameters is a fixed array, not a slice: the size is known at comptime from the
layer list, so there is nothing to allocate and nothing to free.
The test checks two things beyond the count. Every collected Var must be
distinct — a chain handing back the same one twice would train one parameter twice and
another never, which is the same bug arriving by a different route. And the collected list is
compared against attaching the same layers by hand on a fresh graph, so "in order" is asserted
rather than assumed.
znum's Chain iterates a struct's field types and carries a fused-activation tag
table. This one carries a tuple and a count, because the tape fuses nothing and the graph already
composes — there is no second thing for the chain to know.
const h = try bn.attach(&graph, x, .training); // batch stats
const h = try bn.attach(&graph, x, .inference); // running stats
try bn.observe(batch); // move the running statsznum stores mode: Mode on the struct and defaults it to .training,
which is exactly what PyTorch does with model.train() / model.eval(). It
is also one of the most common bugs in either: forget to flip it before evaluating and the
network silently normalises by whatever batch it is looking at.
On a large evaluation batch the numbers look plausible. On a batch of one they are
exactly zero for every feature, because a single sample has no variance — the output
carries no information at all. The model "works in training and breaks in production", and nothing
anywhere reported a mistake. The test makes that visible rather than describing it: it pushes one
row through .training and asserts every output is erased, then pushes the same row
through .inference and asserts it survives.
Here use is a required argument. There is no state to forget, no
default to inherit, and the call site says which statistics it wants. A caller cannot reach the
single-row disaster by omission — only by asking for it.
observe is separate from attach for the same reason.
attach only reads; znum folds the running-statistics update into the forward
pass, so building the graph mutates the layer, and a caller who recomputes a graph twice ages the
statistics twice. The test asserts exactly that: two attaches leave the running mean at zero, and
only observe moves it.
Layer norm normalises each row across its features; batch norm normalises each
column across the batch. That is the same arithmetic on the transpose — verified to
2.2e-16 against a direct per-column implementation — so the forward pass
and its gradient both come from the pair that already exists. znum writes both out
separately; here attach transposes, calls layerNormRows, and transposes
back. Both transposes are views, so the cost is strided access rather than a copy.
const cell = try zn.LstmCell(f64).init(gpa, rng, input_size, hidden_size);
const s = try cell.step(&graph, x, hidden, cell_state);
// s.hidden, s.cell, s.parametersLstmCell/// One LSTM timestep: four gates deciding what to keep, what to add, and what to show.
///
/// THE FORGET BIAS STARTS AT ONE, AND THAT IS THE DIFFERENCE THAT MATTERS
///
/// znum zeroes every bias, and so does PyTorch. A zero forget bias means `sigmoid(0) = 0.5` on
/// the first step, so **the cell halves its memory every timestep before it has learned
/// anything**: after ten steps 0.1% of the original signal survives, and the gradient that would
/// teach it to remember has to travel back through all that decay.
///
/// Starting the forget bias at one makes `sigmoid(1) = 0.73` instead. Measured over ten steps
/// with no input at all, the cell retains **4.3% against 0.1% - forty times more** - which is the
/// difference between a gradient that arrives and one that does not. The result is Gers's, from
/// 1999, and it is still absent from most implementations.
///
/// Every other bias starts at zero, for the usual reason: a random or nonzero bias adds nothing a
/// weight does not already provide.
///
/// FOUR WEIGHT PAIRS, NOT ONE FUSED MATRIX
///
/// PyTorch fuses the four gates into one `(4*hidden, input)` matrix and slices the result, which
/// is one matmul instead of four. That needs a graph-level SLICE with a backward, and this tape
/// has none - so fusing would mean building tape machinery to save three matmuls on a cell whose
/// cost is dominated by the sequence loop around it. **Four pairs, named for their gates**, and
/// the fusion is a note rather than a half-built abstraction.
pub fn LstmCell(comptime T: type) type {
return struct {
/// How much of the new candidate to let in.
input_gate: Gate,
/// How much of the existing cell to keep. Its bias starts at ONE - see above.
forget_gate: Gate,
/// The new value proposed for the cell, before the input gate scales it.
candidate: Gate,
/// How much of the cell to expose as the hidden state.
output_gate: Gate,
const Self = @This();
/// One gate's two weight matrices and its bias.
///
/// Every gate reads BOTH the current input and the previous hidden state, which is the
/// whole idea: `from_input` is what the world says, `from_hidden` is what the cell
/// remembers, and the gate is their sum.
pub const Gate = struct {
from_input: Tensor(T),
from_hidden: Tensor(T),
bias: Tensor(T),
};
/// What one step produces. **Two states, not one** - `hidden` is what the next layer
/// sees, `cell` is the memory that carries forward, and confusing them is the classic
/// LSTM mistake.
pub const Stepped = struct {
hidden: Var,
cell: Var,
parameters: [12]Var,
};
pub fn init(
gpa: Allocator,
rng: Rng,
input_size: usize,
hidden_size: usize,
) Error!Self {
var made: Self = undefined;
const gates = [_][]const u8{ "input_gate", "forget_gate", "candidate", "output_gate" };
inline for (gates, 0..) |name, i| {
const from_input: Tensor(T) = try Tensor(T).alloc(gpa, &.{ input_size, hidden_size });
const from_hidden: Tensor(T) = try Tensor(T).alloc(gpa, &.{ hidden_size, hidden_size });
const bias: Tensor(T) = try Tensor(T).alloc(gpa, &.{ 1, hidden_size });
try initXavier(T, from_input, rng.split(@intCast(i * 2)), input_size, hidden_size);
try initXavier(T, from_hidden, rng.split(@intCast(i * 2 + 1)), hidden_size, hidden_size);
// The forget gate is index 1, and it is the one that starts at one.
bias.fill(if (i == 1) 1 else 0);
@field(made, name) = .{
.from_input = from_input,
.from_hidden = from_hidden,
.bias = bias,
};
}
return made;
}
/// Advance one timestep. Built entirely from tape operations, so the gradient is the
/// tape's and nothing here needs a backward of its own.
pub fn step(
self: Self,
graph: *Graph(T),
input: Var,
hidden: Var,
cell: Var,
) Error!Stepped {
var collected: [12]Var = undefined;
var next: usize = 0;
const gates = [_][]const u8{ "input_gate", "forget_gate", "candidate", "output_gate" };
var summed: [4]Var = undefined;
inline for (gates, 0..) |name, i| {
const gate: Gate = @field(self, name);
const from_input: Var = try graph.parameter(gate.from_input);
const from_hidden: Var = try graph.parameter(gate.from_hidden);
const bias: Var = try graph.parameter(gate.bias);
collected[next] = from_input;
collected[next + 1] = from_hidden;
collected[next + 2] = bias;
next += 3;
const world: Var = try graph.matmul(input, from_input);
const memory: Var = try graph.matmul(hidden, from_hidden);
summed[i] = try graph.add(try graph.add(world, memory), bias);
}
const keep: Var = try graph.sigmoid(summed[1]);
const admit: Var = try graph.sigmoid(summed[0]);
const proposed: Var = try graph.tanh(summed[2]);
const show: Var = try graph.sigmoid(summed[3]);
// The cell is what survives plus what is admitted; the hidden state is the cell seen
// through the output gate. Keeping these two lines apart is the point of the type.
const next_cell: Var = try graph.add(
try graph.mul(keep, cell),
try graph.mul(admit, proposed),
);
const next_hidden: Var = try graph.mul(show, try graph.tanh(next_cell));
return .{ .hidden = next_hidden, .cell = next_cell, .parameters = collected };
}
};
}
The forget bias starts at one. znum zeroes every bias, and so does PyTorch —
and a zero forget bias means sigmoid(0) = 0.5, so the cell halves its memory every
timestep before it has learned anything. Ten steps in, 0.1% of the original signal survives, and
the gradient that would teach it to remember has to travel back through all that decay.
Starting at one gives sigmoid(1) = 0.73. with every weight zeroed so only
the biases decide and the input gate shut so nothing but decay happens:
forget bias 0 0.000977 (0.5^10, as it must be)
forget bias 1 0.043604 (0.73^10)Forty-four times the signal, and with it forty-four times the gradient. The
result is Gers's, from 1999, and it is still absent from most implementations. The test asserts the
exact powers rather than just the ratio, so it measures the gate and not the arithmetic
around it — and it separately asserts that init produces this default, because a
default nobody checks is a default that drifts.
step returns hidden and cell. The hidden state is
what the next layer sees; the cell is the memory that carries forward. A type returning one value
could not stop you mixing them, so this one returns both, and they are different
vars.
PyTorch fuses the four gates into one (4·hidden, input) matrix and slices the
result — one matmul instead of four. That needs a graph-level slice with a backward,
and this tape has none. Fusing would mean building tape machinery to save three matmuls on a cell
whose cost is dominated by the sequence loop around it, so the fusion is a note here rather than a
half-built abstraction.
znum keeps twelve separately named tensors (w_ii, w_hi,
b_i, and so on). Here they are grouped into four Gates with
from_input and from_hidden, because that grouping is the idea:
every gate reads both what the world says and what the cell remembers.
Every operation in step is already on the tape — matmul, sigmoid, tanh, add,
mul — so the gradient is the tape's and nothing here needs a backward of its own.
try attention.attach(&graph, x, .everywhere); // encoder: all positions
try attention.attach(&graph, x, .backward_only); // decoder: no peeking aheadSame shape as BatchNorm.Use, and the consequence is worse. A language model trained
without a causal mask can see the answer it is predicting: the loss drops to near
zero, the training curve looks superb, and the model generates nonsense the moment it has to
predict something it has not been shown. Nothing fails, nothing warns, and the metric that
would tell you is the one that looks best.
So look is required at every call. There is no default to inherit and no field to
forget.
Inspecting the mask would only confirm the mask is the one intended. Instead, replacing the last position with a large value and measures how much every earlier output moves:
.backward_only largest shift EXACTLY 0
.everywhere largest shift > 0.1Not "small" — zero. A causal mask that let through a millionth would be a causal mask with a bug, and a tolerance would hide it. The unmasked run is what proves the zero came from the mask rather than from the input being ignored.
The scale is derived, not passed. Dot products grow with the width they are
summed over, so without 1/sqrt(width) the softmax saturates and the gradient vanishes.
Passing the wrong dimension is a quiet, common mistake; here it comes from the projection's own
shape and cannot disagree with the tensors it scales.
The mask adds a large negative, not -inf. Negative infinity is
exactly right in algebra and produces inf - inf = NaN the moment a row is entirely
masked — which the first row of a causal mask never is, but a padded batch would be. A number
the exponential flushes to zero says the same thing and survives.
try zn.Attention(f64).init(gpa, rng, model_dim, width, heads);Attention/// Scaled dot-product attention: every position reads every position it is allowed to see.
///
/// WHETHER IT CAN SEE THE FUTURE IS AN ARGUMENT, NOT A SETTING
///
/// The same lesson as `BatchNorm.Use`, and the consequence is worse. A language model trained
/// without a causal mask **can see the answer it is predicting**: the loss drops to near zero,
/// the training curve looks superb, and the model generates nonsense the moment it has to predict
/// something it has not been shown. Nothing fails, nothing warns, and the metric that would tell
/// you is the one that looks best.
///
/// So `look` is required at every call:
///
/// try attention.attach(&graph, x, .everywhere); // encoder: all positions
/// try attention.attach(&graph, x, .backward_only); // decoder: no peeking ahead
///
/// THE SCALE IS DERIVED, NOT PASSED
///
/// Dot products grow with the width they are summed over, so without `1/sqrt(width)` the softmax
/// saturates and the gradient vanishes - and passing the wrong dimension is a quiet, common
/// mistake. Here it comes from the projection's own shape, so it cannot disagree with the tensors
/// it scales.
///
/// SINGLE HEAD. Multiple heads need the projection split and the results concatenated, and this
/// tape has neither a slice nor a concat with a backward - the same limit that stopped `LstmCell`
/// fusing its gates. Stated rather than half-built.
pub fn Attention(comptime T: type) type {
return struct {
/// What each position is looking for.
query: Tensor(T),
/// What each position offers as a label.
key: Tensor(T),
/// What each position hands over when chosen.
value: Tensor(T),
/// How many independent attention patterns share the projection width.
heads: usize,
const Self = @This();
/// Which positions a query may attend to. **Required** - see the note above.
pub const Look = enum {
/// Every position, in both directions. What an encoder wants.
everywhere,
/// Only positions at or before this one. What a decoder MUST have.
backward_only,
};
pub const Attached = struct {
out: Var,
query: Var,
key: Var,
value: Var,
};
/// `heads` independent attention patterns sharing one projection of `width` columns.
///
/// `heads = 1` is ordinary attention. More heads do not widen anything - they DIVIDE the
/// same width, so each head attends in a narrower subspace and the cost is unchanged.
/// That is the part the name hides: multi-head is not more capacity, it is the same
/// capacity spent on several patterns instead of one.
pub fn init(
gpa: Allocator,
rng: Rng,
model_dim: usize,
width: usize,
heads: usize,
) Error!Self {
if (heads == 0 or width % heads != 0) {
return Error.UnsupportedShape;
}
const query: Tensor(T) = try Tensor(T).alloc(gpa, &.{ model_dim, width });
const key: Tensor(T) = try Tensor(T).alloc(gpa, &.{ model_dim, width });
const value: Tensor(T) = try Tensor(T).alloc(gpa, &.{ model_dim, width });
try initXavier(T, query, rng.split(0), model_dim, width);
try initXavier(T, key, rng.split(1), model_dim, width);
try initXavier(T, value, rng.split(2), model_dim, width);
return .{ .query = query, .key = key, .value = value, .heads = heads };
}
pub fn attach(self: Self, graph: *Graph(T), input: Var, look: Look) Error!Attached {
const x: Tensor(T) = graph.valueOf(input);
if (x.rank != 2 or x.shape[1] != self.query.shape[0]) {
return Error.ShapeMismatch;
}
const positions: usize = x.shape[0];
const width: usize = self.query.shape[1];
const wq: Var = try graph.parameter(self.query);
const wk: Var = try graph.parameter(self.key);
const wv: Var = try graph.parameter(self.value);
const q: Var = try graph.matmul(input, wq);
const k: Var = try graph.matmul(input, wk);
const v: Var = try graph.matmul(input, wv);
// RANK 2 THROUGHOUT, ONE HEAD AT A TIME
//
// znum splits heads by reshaping to rank 4 and permuting: `[rows, D]` becomes
// `[B, S, H, d]` becomes `[B, H, S, d]`. That needs a differentiable reshape AND
// permute, and it makes every line afterwards reason about four axes.
//
// With `graph.slice` and `graph.concat` the same thing is a loop: take each head's
// columns, attend, and join the results. Every tensor stays rank 2 and each iteration
// reads exactly like the single-head maths, because it IS the single-head maths.
//
// The cost: `heads` separate matmuls where a batched form would do one. For the head
// counts a rank-2 tape sees that is a fair trade for not having two more tape ops and
// four axes to track - and it is stated here rather than discovered.
const head_width: usize = width / self.heads;
const scale_by: T = 1.0 / @sqrt(@as(T, @floatFromInt(head_width)));
var joined: ?Var = null;
for (0..self.heads) |head| {
const from: usize = head * head_width;
const qh: Var = try graph.slice(q, 1, from, head_width);
const kh: Var = try graph.slice(k, 1, from, head_width);
const vh: Var = try graph.slice(v, 1, from, head_width);
const raw_head: Var = try graph.matmul(qh, try graph.transpose(kh));
const scaled_head: Var = try graph.scale(raw_head, scale_by);
const allowed: Var = try applyLook(graph, scaled_head, look, positions);
const attended: Var = try graph.matmul(try graph.softmaxRows(allowed), vh);
joined = if (joined) |so_far|
try graph.concat(so_far, attended, 1)
else
attended;
}
return .{ .out = joined.?, .query = wq, .key = wk, .value = wv };
}
/// The mask, or nothing, depending on `look`.
fn applyLook(graph: *Graph(T), scores: Var, look: Look, positions: usize) Error!Var {
if (look == .everywhere) {
return scores;
}
// A LARGE NEGATIVE, NOT AN INFINITY. `-inf` would be exactly right in algebra and
// produces `inf - inf = NaN` the moment a row is entirely masked, which the first row
// of a causal mask never is but a padded batch would be. A number the exponential
// flushes to zero says the same thing and survives.
const mask: Tensor(T) = try Tensor(T).alloc(graph.gpa, &.{ positions, positions });
for (0..positions) |row| {
for (0..positions) |col| {
try mask.setAt(&.{ row, col }, if (col > row) -1.0e30 else 0);
}
}
return graph.add(scores, try graph.constant(mask));
}
};
}
More heads do not widen anything — they divide the same width. Two heads over a width of 8 produce exactly the same output shape and the same parameter count as one head over 8; each simply attends in a narrower subspace. That is the part "multi-head" hides, and this is exact for it: same shape, different numbers.
The scale then comes from the head's width, not the projection's.
Splitting 8 into two heads means each sums over 4, so the scale is 1/sqrt(4). Getting
that wrong is invisible in the shapes and quiet in the output — it just saturates the softmax
a little more than it should. The test pins it by checking that one head over 4 and two heads over
8 share a scale, where a projection-wide scale would have made them differ.
A width that does not divide by the head count is refused rather than silently dropping a column.
znum splits heads by reshaping to rank 4 and permuting: [rows, D] becomes
[B, S, H, d] becomes [B, H, S, d]. That needs a differentiable reshape
and permute, and every line afterwards reasons about four axes.
With graph.slice and graph.concat the same thing is a loop: take each
head's columns, attend, join the results. Every tensor stays rank 2, and each iteration reads
exactly like the single-head maths — because it is the single-head maths.
The cost is heads separate matmuls where a batched form would do one. For the head
counts a rank-2 tape sees, that is a fair trade for not carrying two more tape ops and four axes,
and it is stated here rather than discovered.
h = x + attention(norm(x))
y = h + feedForward(norm(h))The block normalises before each sublayer and adds the result to the untouched input. The original transformer did it the other way — sublayer, add, then normalise — and that version cannot be trained without a learning-rate warmup, because the normalisation sits on the residual path and rescales the gradient at every block.
Pre-norm leaves the path clear: the gradient reaches the input carrying a component of exactly one from each residual, whatever the depth. znum is pre-norm too; the reason usually is not written down, so it is here — and the consequence rather than the arrangement, by measuring that no element of the input gradient has collapsed.
A fresh block is close to the identity. The sublayers start small against the residual, so depth you have not trained yet does nothing rather than something arbitrary. That is what makes stacking blocks safe before any training has happened, and there is a bound on how far the output can move from its input.
The mask lives inside Attention. But a residual carries the input forward around
the attention, so composition is exactly where a leak would appear — the block could route
the future past its own mask. The same perturbation test therefore runs again at this level:
disturb the last position, and every earlier output must be bit-identical.
The seven parameters are collected the same way Chain does it, for the same reason:
one missing from that list would silently never train.
try zn.TransformerBlock(f64).init(gpa, rng, model_dim, hidden, heads, eps, .swiglu);TransformerBlock/// One transformer block: attention and a feed-forward pair, each behind a residual.
///
/// PRE-NORM, AND THE RESIDUAL IS WHY
///
/// The block normalises BEFORE each sublayer and adds the result to the untouched input:
///
/// h = x + attention(norm(x))
/// y = h + feedForward(norm(h))
///
/// The original transformer did it the other way - sublayer, add, then normalise - and that
/// version **cannot be trained without a learning-rate warmup**, because the normalisation sits
/// on the residual path and rescales the gradient every block. Pre-norm leaves the path clear:
/// the gradient reaches the input with a component of exactly one from each residual, whatever
/// the depth. znum is pre-norm too; this note exists because the reason usually is not written
/// down, and the test below asserts the consequence rather than the arrangement.
///
/// The sublayers begin small relative to the residual, so a fresh block is close to the identity.
/// That is the point: depth you have not trained yet should do nothing, not something arbitrary.
pub fn TransformerBlock(comptime T: type) type {
return struct {
/// Reads across positions.
attention: Attention(T),
/// The first half of the position-wise pair, widening.
widen: Dense(T),
/// The second half, narrowing back so the residual still fits.
narrow: Dense(T),
/// Added under the square root in both norms.
epsilon: T,
/// Which feed-forward this block was built with - it decided `widen`'s shape.
feed_forward: FeedForward,
const Self = @This();
pub const Attached = struct {
out: Var,
/// Everything the block attached, ready for the optimiser - the same argument as
/// `Chain`: a parameter missing from this list would silently never train.
parameters: [7]Var,
};
/// Which feed-forward the block uses. **Chosen at `init`, not at `attach`.**
///
/// THE GATE CHANGES A LAYER'S SHAPE, NOT JUST THE ACTIVATION BETWEEN TWO
///
/// It looks like an activation swap and is not. SwiGLU consumes half its input as a gate,
/// so to produce `hidden` columns the widening layer has to make **2 x hidden** - which
/// makes it a different network, not the same one with a different nonlinearity.
///
/// relu widen (model -> hidden), relu, narrow (hidden -> model)
/// swiglu widen (model -> 2 x hidden), swiglu, narrow (hidden -> model)
///
/// At `model_dim = 4, hidden = 16` that is 148 parameters against 228. A caller comparing
/// the two is comparing capacities as well as gates, and the doc says so rather than
/// letting them conclude the gate alone did it.
pub const FeedForward = enum { relu, swiglu };
/// `hidden` is the width of the feed-forward middle. Four times `model_dim` is the usual
/// choice and is the caller's to make rather than a constant hidden here.
pub fn init(
gpa: Allocator,
rng: Rng,
model_dim: usize,
hidden: usize,
heads: usize,
epsilon: T,
feed_forward: FeedForward,
) Error!Self {
// SwiGLU splits the widened columns in half, so it needs twice as many to hand the
// same width to the narrowing layer.
const widened: usize = switch (feed_forward) {
.relu => hidden,
.swiglu => hidden * 2,
};
return .{
.attention = try Attention(T).init(gpa, rng.split(0), model_dim, model_dim, heads),
.widen = try Dense(T).init(gpa, rng.split(1), model_dim, widened, .xavier),
.narrow = try Dense(T).init(gpa, rng.split(2), hidden, model_dim, .xavier),
.epsilon = epsilon,
.feed_forward = feed_forward,
};
}
pub fn attach(self: Self, graph: *Graph(T), input: Var, look: Attention(T).Look) Error!Attached {
// Attention sublayer. `input` goes into the add untouched - that is the residual.
const before_attention: Var = try graph.layerNormRows(input, self.epsilon);
const attended = try self.attention.attach(graph, before_attention, look);
const mixed: Var = try graph.add(input, attended.out);
// Feed-forward sublayer, position by position, with its own residual.
const before_feedforward: Var = try graph.layerNormRows(mixed, self.epsilon);
const wide = try self.widen.attach(graph, before_feedforward);
const activated: Var = switch (self.feed_forward) {
.relu => try graph.relu(wide.out),
// Halves the width, which is why `init` doubled it.
.swiglu => try graph.swiglu(wide.out),
};
const narrowed = try self.narrow.attach(graph, activated);
return .{
.out = try graph.add(mixed, narrowed.out),
.parameters = .{
attended.query, attended.key, attended.value,
wide.weight, wide.bias, narrowed.weight,
narrowed.bias,
},
};
}
};
}
It looks like an activation swap and is not. SwiGLU consumes half its input as a
gate, so to hand hidden columns to the narrowing layer, the widening layer has
to produce 2 x hidden:
relu widen (model -> hidden), relu, narrow (hidden -> model)
swiglu widen (model -> 2 x hidden), swiglu, narrow (hidden -> model)At model_dim = 4, hidden = 16 that is 148 parameters against 228.
A caller comparing the two is comparing capacities as well as gates, which is worth knowing before
concluding the gate alone moved a training curve. The test asserts both widths and that the
narrowing layer is unchanged.
What does not change: the output shape, so a block is still a block and still stacks; and the parameter count, still seven — the gate changed a shape, not the number of things to train.
Causality is re-checked for the gated block too. The mask is in the attention, but a residual carries the input around it either way, so the perturbation test runs again rather than being assumed to transfer.
const h = try graph.glu(x); // slice * sigmoid(slice)
const h = try graph.swiglu(x); // slice * silu(slice)Half the columns carry the value and half the gate, so the output is half as wide. SwiGLU is the gate in every recent transformer feed-forward.
Neither has a backward. znum writes gatedForward and then a
hand-derived gatedBackward — walking coordinates and applying the product rule by
hand. Here a GLU is
slice(x, 0, half) * sigmoid(slice(x, half, half))and every one of those is already a tape node, so the gradient is the tape's. There is no second derivation to keep in step with the first, and no chance of the pair disagreeing under maintenance.
This is the third thing graph.slice and graph.concat paid for, after
multi-head attention and the batched cartpole twin. They were built because their absence
had shaped two designs; they keep shaping the ones that come after.
The test still checks the gradient against central differences at every element — over the value half and the gate half, since a product rule applied by hand is exactly the kind of thing that is right on one branch and wrong on the other. And it asserts SwiGLU differs numerically from GLU, because a version ignoring the gate kind would pass the shape check and every gradient check above it.
try zn.saveParameters(f64, gpa, &bytes, &graph, run.parameters);
// ... later, in a fresh process ...
try zn.loadParameters(f64, gpa, bytes, &graph, run.parameters);saveParameterspub fn saveParameters(
comptime T: type,
gpa: Allocator,
out: *std.ArrayList(u8),
graph: *const Graph(T),
parameters: []const Var,
) Error!void {
comptime requireFloat(T);
const names: [][]const u8 = gpa.alloc([]const u8, parameters.len) catch
return Error.OutOfMemory;
defer {
for (names) |name| {
gpa.free(name);
}
gpa.free(names);
}
const tensors: []Tensor(T) = gpa.alloc(Tensor(T), parameters.len) catch
return Error.OutOfMemory;
defer gpa.free(tensors);
for (parameters, 0..) |parameter, i| {
names[i] = std.fmt.allocPrint(gpa, "p{d}", .{i}) catch return Error.OutOfMemory;
tensors[i] = graph.valueOf(parameter);
}
return saveTensors(T, gpa, out, names, tensors);
}
loadParameters/// Read a network's parameters back into it, in place.
///
/// The tensors are written THROUGH the graph's parameter values, so a network built with the
/// same architecture picks up the saved weights and is ready to run. `loadTensors` refuses a
/// shape that disagrees, which catches the common mistake of loading into a different
/// architecture.
pub fn loadParameters(
comptime T: type,
gpa: Allocator,
bytes: []const u8,
graph: *const Graph(T),
parameters: []const Var,
) Error!void {
comptime requireFloat(T);
const names: [][]const u8 = gpa.alloc([]const u8, parameters.len) catch
return Error.OutOfMemory;
defer {
for (names) |name| {
gpa.free(name);
}
gpa.free(names);
}
const tensors: []Tensor(T) = gpa.alloc(Tensor(T), parameters.len) catch
return Error.OutOfMemory;
defer gpa.free(tensors);
for (parameters, 0..) |parameter, i| {
names[i] = std.fmt.allocPrint(gpa, "p{d}", .{i}) catch return Error.OutOfMemory;
tensors[i] = graph.valueOf(parameter);
}
return loadTensors(T, bytes, names, tensors);
}
saveTensors and loadTensors were always here, and they take a list of
names and a list of tensors. What was missing is the bridge from a network to those two
lists — so this library could train a cartpole agent and not keep it, and every run started
from noise.
Not that the bytes round-trip — the tensor tests already make that claim. The claim is that a network rebuilt from a different seed and then loaded produces the same output as the one that trained.
Three things make that meaningful rather than decorative:
The parameters are saved as p0, p1, ... in attach order,
because a Var carries no name. The loading network must have the same
architecture in the same order.
loadTensors checks the shapes, so a changed layer is caught — loading
a checkpoint into a network whose hidden width grew and gets a ShapeMismatch. A
same-shaped layer in a different role is not caught, and that is worth knowing rather than
discovering.
const norm = try zn.clipGradNorm(f64, &graph, run.parameters, 1.0);
// `norm` is the value BEFORE clipping - log itclipGradNorm/// Scale every gradient down together if their combined norm exceeds `max_norm`.
///
/// Returns the norm BEFORE clipping, which is the number worth logging: a spike in it is how you
/// find out a batch was bad, and clipping silently would hide exactly that.
///
/// ONE NORM ACROSS ALL PARAMETERS, NOT ONE PER TENSOR
///
/// This is the part that is easy to get wrong and hard to notice. Clipping each tensor to its own
/// limit **changes the direction of the update**: a layer whose gradient was twice its neighbour's
/// comes out the same size, so the step now points somewhere the loss surface never suggested.
///
/// Scaling everything by one factor keeps the direction exactly and shortens the step, which is
/// what "clip the gradient" is supposed to mean. The test asserts the direction survives, because
/// a per-tensor version passes every test about magnitude.
///
/// WHY THIS MATTERS MORE THAN THE OPTIMISER YOU PICK
///
/// A transformer meets a bad batch, the loss spikes, and one enormous update moves the weights
/// somewhere the model never recovers from. Adam does not save you - it normalises by a RUNNING
/// average, so a single outlier is still an outlier. Clipping is the thing that does.
pub fn clipGradNorm(
comptime T: type,
graph: *const Graph(T),
parameters: []const Var,
max_norm: T,
) Error!T {
comptime requireFloat(T);
if (max_norm <= 0) {
return Error.DomainError;
}
var squares: CompensatedSum(T) = .{};
for (parameters) |parameter| {
const slopes: Tensor(T) = try graph.gradOf(parameter);
for (0..slopes.size()) |i| {
const g: T = slopes.data[i];
squares.add(g * g);
}
}
const total: T = @sqrt(squares.value());
if (total <= max_norm) {
return total;
}
// A SINGLE FACTOR, APPLIED TO EVERYTHING - that is what keeps the direction.
const factor: T = max_norm / total;
for (parameters) |parameter| {
const slopes: Tensor(T) = try graph.gradOf(parameter);
for (0..slopes.size()) |i| {
slopes.data[i] *= factor;
}
}
return total;
}
A transformer meets a bad batch, the loss spikes, and one enormous update moves the weights somewhere the model never recovers from. Adam does not save you — it normalises by a running average, so a single outlier is still an outlier. Clipping is the thing that does.
The return value is the norm before clipping, which is the number worth logging: a spike in it is how you learn a batch was bad, and clipping silently would hide exactly that.
This is the part that is easy to get wrong and hard to notice. Say two parameters whose gradients have norms 5 and 1:
global both scaled by 1/sqrt(26) -> norms 0.98 and 0.196, ratio still 5
per-tensor each scaled to 1 -> norms 1 and 1, ratio now 1Clipping each tensor to its own limit changes the direction of the update. The layer whose gradient was five times its neighbour's comes out the same size, and the step now points somewhere the loss surface never suggested.
Scaling everything by one factor keeps the direction exactly and shortens the step, which is what "clip the gradient" is supposed to mean.
The test asserts the ratio survives, not just the final magnitude — a per-tensor version passes every check about size. Verified by writing the per-tensor version and watching the test fail.
One asymmetry worth noting: a gradient already under the limit is left alone, not scaled up to meet it. That would be a different operation entirely, and a tempting symmetry to add.
const v = try zn.variance(f64, t, .sample);
const sd = try zn.stdDev(f64, t, .population);
const r = try zn.rootMeanSquare(f64, t);
const d = try zn.meanAbsDev(f64, t);
const se = try zn.stdErr(f64, t);
try zn.zscore(f64, out, t); // mean 0, sample deviation 1
try zn.minMaxScale(f64, out, t); // into [0, 1]
const c = try zn.covariance(f64, a, b, .sample);
const corr = try zn.correlation(f64, a, b);variance/// The variance of every element, with the divisor chosen by `ddof`.
///
/// TWO PASSES, FOR THE REASON `layerNormRows` GIVES
///
/// `E[x^2] - E[x]^2` is one pass and subtracts two large nearly-equal numbers, losing most of its
/// significant digits when the mean is large relative to the spread. Two passes cost one more
/// read of memory that is already in cache and are correct for every input.
///
/// Returns `DomainError` when the divisor is undefined - `n = 0`, or `n = 1` with
/// `.sample` - rather than a NaN that propagates into a result nobody examines. That is what
/// `Ddof.divisor` returning an optional is for.
pub fn variance(comptime T: type, a: Tensor(T), ddof: Ddof) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
const divisor: f64 = ddof.divisor(n) orelse return Error.DomainError;
const mean: T = try meanAll(T, a);
var total: CompensatedSum(T) = .{};
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
const d: T = (try a.at(at)) - mean;
const term: T = d * d;
total.add(term);
}
const divide_by: T = @floatCast(divisor);
return total.value() / divide_by;
}
stdDev/// The standard deviation: the square root of `variance`.
pub fn stdDev(comptime T: type, a: Tensor(T), ddof: Ddof) Error!T {
return @sqrt(try variance(T, a, ddof));
}
rootMeanSquare/// The root mean square of every element.
pub fn rootMeanSquare(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
const x: T = try a.at(at);
const term: T = x * x;
total.add(term);
}
return @sqrt(total.value() / @as(T, @floatFromInt(n)));
}
meanAbsDev/// The mean absolute deviation from the mean.
pub fn meanAbsDev(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n == 0) {
return Error.DomainError;
}
const mean: T = try meanAll(T, a);
var total: T = 0;
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
total += @abs((try a.at(at)) - mean);
}
return total / @as(T, @floatFromInt(n));
}
stdErr/// The standard error of the mean: `stdDev(.sample) / sqrt(n)`.
pub fn stdErr(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n < 2) {
return Error.DomainError;
}
return (try stdDev(T, a, .sample)) / @sqrt(@as(T, @floatFromInt(n)));
}
zscore/// `out = (a - mean) / stdDev`, elementwise - the z-score of every element.
///
/// Uses the SAMPLE deviation, matching what a reader gets from R or scipy by default. The
/// population form is a different number and the difference is visible at small n, so the choice
/// is stated rather than left to whichever the implementer reached for.
pub fn zscore(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireFloat(T);
const mean: T = try meanAll(T, a);
const sd: T = try stdDev(T, a, .sample);
if (sd == 0) {
return Error.DomainError;
}
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const va: Tensor(T) = try a.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const x: T = va.data[va.offsetOf(at)];
out.data[out.offsetOf(at)] = (x - mean) / sd;
}
}
minMaxScale/// Normalise every element into `[0, 1]` by the tensor's own range.
pub fn minMaxScale(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void {
comptime requireNumeric(T);
const lo: T = try minAll(T, a);
const hi: T = try maxAll(T, a);
if (hi == lo) {
return Error.DomainError;
}
if (out.isAliased()) {
return Error.UnsupportedShape;
}
const shape: []const usize = out.shape[0..out.rank];
const va: Tensor(T) = try a.broadcastTo(shape);
var walk: Walk = .over(shape);
while (walk.next()) |at| {
const x: T = va.data[va.offsetOf(at)];
out.data[out.offsetOf(at)] = (x - lo) / (hi - lo);
}
}
covariance/// The covariance of two tensors of equal size, with the divisor chosen by `ddof`.
pub fn covariance(comptime T: type, a: Tensor(T), b: Tensor(T), ddof: Ddof) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n != b.size()) {
return Error.ShapeMismatch;
}
const divisor: f64 = ddof.divisor(n) orelse return Error.DomainError;
const mean_a: T = try meanAll(T, a);
const mean_b: T = try meanAll(T, b);
var total: CompensatedSum(T) = .{};
var pair: Walk2 = .over(a.shape[0..a.rank], b.shape[0..b.rank]);
while (pair.next()) |at| {
const x: T = (try a.at(at.first)) - mean_a;
const y: T = (try b.at(at.second)) - mean_b;
const term: T = x * y;
total.add(term);
}
const divide_by: T = @floatCast(divisor);
return total.value() / divide_by;
}
correlation/// Pearson's correlation coefficient of two tensors of equal size.
///
/// NOT clamped to [-1, 1]. Rounding can put it a few ULP outside, and clamping would hide the
/// only signal that something upstream is wrong - a correlation of 1.0000003 says the inputs are
/// collinear AND that the arithmetic is at its limit; a clamped 1.0 says only the first.
pub fn correlation(comptime T: type, a: Tensor(T), b: Tensor(T)) Error!T {
comptime requireFloat(T);
const sa: T = try stdDev(T, a, .sample);
const sb: T = try stdDev(T, b, .sample);
if (sa == 0 or sb == 0) {
return Error.DomainError;
}
return (try covariance(T, a, b, .sample)) / (sa * sb);
}
variance is two-pass, for layerNormRows's reason. Both it and
covariance take a Ddof and return DomainError where the
divisor is undefined. zscore and stdErr use the sample deviation,
matching R and scipy.
correlation is not clamped to [−1, 1]. Rounding can put it a few ULP
outside, and a value slightly beyond the range is a signal about the inputs' conditioning that
clamping would erase.
minMaxScale and zscore return DomainError for a constant
tensor, which has no range and no deviation to divide by.
const q = try zn.quantile(f64, a, 0.9, scratch); // scratch is sorted in place
const mid = try zn.median(f64, a, scratch);
const lean = try zn.skew(f64, a);
const tails = try zn.kurtosis(f64, a); // excess: normal reads as zero
const g = try zn.geoMean(f64, a);
const hm = try zn.harmMean(f64, a);
try zn.histogram(f64, &counts, a, low, high);quantile/// The value below which a fraction `q` of the data lies, by linear interpolation between the
/// two neighbouring order statistics.
///
/// `scratch` must be at least `a.size()` long and is overwritten with the sorted values, so a
/// caller wanting several quantiles sorts once and calls `quantileSorted` for each.
///
/// THE INTERPOLATION RULE IS A CHOICE, AND THIS IS THE COMMON ONE
///
/// `q` maps to position `q * (n - 1)` in the sorted array: q=0 is the first element, q=1 the
/// last, and anything between is a weighted blend of its two neighbours. Nine such rules are in
/// common use and they disagree on small samples; this is the one numpy calls "linear" and the
/// one znum uses, so a reader comparing against either gets the same number.
pub fn quantile(comptime T: type, a: Tensor(T), q: T, scratch: []T) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n == 0 or scratch.len < n) {
return Error.ShapeMismatch;
}
if (q < 0 or q > 1) {
return Error.DomainError;
}
var walker: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
for (0..n) |i| {
scratch[i] = a.data[a.offsetOf(walker[0..a.rank])];
advance(shape, walker[0..a.rank]);
}
std.mem.sort(T, scratch[0..n], {}, std.sort.asc(T));
return quantileSorted(T, scratch[0..n], q);
}
median/// The middle value: `quantile(0.5)`.
pub fn median(comptime T: type, a: Tensor(T), scratch: []T) Error!T {
comptime requireFloat(T);
return quantile(T, a, 0.5, scratch);
}
skew/// The third standardised moment: how far the distribution leans.
///
/// Zero for anything symmetric, positive when the long tail is to the right. This is the
/// POPULATION form, dividing by `n` and by the population deviation cubed - the same as znum's,
/// and the same as what most plotting libraries show. The sample-corrected form multiplies by
/// `sqrt(n(n-1))/(n-2)` and matters below about twenty points.
pub fn skew(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
return standardisedMoment(T, a, 3);
}
kurtosis/// The fourth standardised moment, minus three: how heavy the tails are.
///
/// Three is subtracted so that a normal distribution reads as ZERO rather than as three, which
/// is what "excess kurtosis" means and what every reader expects when a number is called
/// kurtosis. Positive is heavier-tailed than normal.
pub fn kurtosis(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
return (try standardisedMoment(T, a, 4)) - 3;
}
geoMean/// The geometric mean: `exp(mean(log x))`. Every value must be positive.
///
/// Written through logarithms rather than as the n-th root of the product, because the product
/// of a few hundred values well inside the float range still overflows it. The logarithm turns
/// the multiplication into an addition that cannot.
pub fn geoMean(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
const x: T = try a.at(at);
if (x <= 0) {
return Error.DomainError;
}
total.add(@log(x));
}
return @exp(total.value() / @as(T, @floatFromInt(n)));
}
harmMean/// The harmonic mean: `n / sum(1/x)`. Every value must be positive.
///
/// The right average for rates: two legs of a journey at 30 and 60 give a harmonic mean of 40,
/// which is the speed that covers the same distance in the same time. The arithmetic mean of 45
/// answers a question nobody asked.
pub fn harmMean(comptime T: type, a: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n == 0) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
const x: T = try a.at(at);
if (x <= 0) {
return Error.DomainError;
}
total.add(1 / x);
}
return @as(T, @floatFromInt(n)) / total.value();
}
histogram/// Count values into `counts.len` equal-width bins spanning `[low, high)`, with the top edge
/// closed so the largest value lands in the last bin rather than outside every bin.
///
/// Values outside the range are not counted, and `Error.DomainError` if the range is empty.
pub fn histogram(
comptime T: type,
counts: []usize,
a: Tensor(T),
low: T,
high: T,
) Error!void {
comptime requireFloat(T);
if (counts.len == 0 or !(high > low)) {
return Error.DomainError;
}
@memset(counts, 0);
const bins: T = @floatFromInt(counts.len);
const width: T = (high - low) / bins;
var walker: [max_rank]usize = @splat(0);
const shape: []const usize = a.shape[0..a.rank];
var remaining: usize = a.size();
while (remaining > 0) : (remaining -= 1) {
const x: T = a.data[a.offsetOf(walker[0..a.rank])];
advance(shape, walker[0..a.rank]);
if (x < low or x > high) {
continue;
}
// The top edge is closed: without this, `high` itself falls one past the last bin.
const slot: usize = if (x == high)
counts.len - 1
else
@intFromFloat(@floor((x - low) / width));
counts[@min(slot, counts.len - 1)] += 1;
}
}
Nine interpolation rules for quantiles are in common use and they disagree on small samples.
This is the one numpy calls "linear": q maps to position
q·(n−1), so q=0 is the first element, q=1 the last, and anything between
is a weighted blend of its two neighbours. The test asserts a case that does not land on
an element, because an all-integer sample would hide the rule entirely.
kurtosis subtracts three so a normal distribution reads as zero
rather than as three — that is what "excess kurtosis" means and what a reader expects from a
number with that name. The test pins it against a uniform distribution, whose excess kurtosis is
exactly −6/5: a number that merely looked small would not distinguish a
correct fourth moment from a slightly wrong one.
geoMean goes through logarithms rather than taking the n-th root of a product,
because the product of a few hundred values well inside the float range still overflows it. Three
hundred values of 1e5 have a product near 1e1500; the answer is still exactly 1e5.
The three means are tested against each other: harmonic ≤ geometric ≤ arithmetic holds for any positive sample, with equality only when every value is the same. That one inequality checks all three at once, where three separate expected numbers would check none of them against the others.
const spread = try zn.medianAbsDev(f64, scratch, a);
const relative = try zn.coefficientOfVariation(f64, a, .sample);
const err = try zn.standardError(f64, a, .sample);
const common = try zn.mode(f64, scratch, a);
const n = try zn.valueCounts(f64, &values, &counts, scratch, a);
try zn.correlationMatrix(f64, out, &.{ x, y, z });medianAbsDev/// The median of the absolute deviations from the median.
///
/// TWO FUNCTIONS, NOT ONE WITH A FLAG - BECAUSE "MAD" MEANS BOTH
///
/// `meanAbsDev` already sits beside this one. znum has a single `mad` with a `MadKind` selecting
/// between them, which means every call site reads `mad(x, .median)` and every reader has to
/// know the enum. Two names say it without one.
///
/// The median version is the robust one: half the data can be arbitrarily wrong and this barely
/// budges, where a single outlier can dominate a standard deviation. That is the whole reason it
/// exists, so it is the one whose name should be unambiguous.
pub fn medianAbsDev(comptime T: type, scratch: []T, a: Tensor(T)) Error!T {
comptime requireFloat(T);
if (scratch.len < a.size()) {
return Error.ShapeMismatch;
}
const centre: T = try median(T, a, scratch[0..a.size()]);
var index: usize = 0;
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| : (index += 1) {
scratch[index] = @abs((try a.at(at)) - centre);
}
std.mem.sort(T, scratch[0..a.size()], {}, std.sort.asc(T));
return quantileSorted(T, scratch[0..a.size()], 0.5);
}
coefficientOfVariation/// The standard deviation as a fraction of the mean.
///
/// Spelled out rather than `cv`, because two letters that stand for three words are a lookup.
/// It answers "is this spread large FOR THIS SCALE" - a deviation of 3 means nothing until you
/// know whether the mean is 5 or 5000.
///
/// A mean of zero has no scale to be a fraction of, and that is a `DomainError` rather than an
/// infinity: the question does not apply, which is different from the answer being large.
pub fn coefficientOfVariation(comptime T: type, a: Tensor(T), ddof: Ddof) Error!T {
comptime requireFloat(T);
const centre: T = try meanAll(T, a);
if (centre == 0) {
return Error.DomainError;
}
return @sqrt(try variance(T, a, ddof)) / @abs(centre);
}
standardError/// How far the MEAN of a sample is likely to be from the mean of the whole.
///
/// The standard deviation shrunk by the square root of the count: more data, steadier estimate.
/// Spelled out rather than `sem`, on the same grounds as `coefficientOfVariation`.
pub fn standardError(comptime T: type, a: Tensor(T), ddof: Ddof) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n == 0) {
return Error.DomainError;
}
return @sqrt(try variance(T, a, ddof)) / @sqrt(@as(T, @floatFromInt(n)));
}
mode/// The most common value, ties going to the smallest.
///
/// TIES GO TO THE SMALLEST, AND THAT IS A CHOICE THE CALLER SHOULD NOT HAVE TO GUESS
///
/// Two values can be equally common, and "the mode" is then ambiguous. Returning whichever the
/// scan happened to see first makes the answer depend on element order, which is exactly the
/// kind of thing that is reproducible until the day it is not. The smallest is arbitrary too,
/// but it is arbitrary the SAME WAY every time.
pub fn mode(comptime T: type, scratch: []T, a: Tensor(T)) Error!T {
comptime requireFloat(T);
const n: usize = a.size();
if (n == 0) {
return Error.DomainError;
}
if (scratch.len < n) {
return Error.ShapeMismatch;
}
try sort(T, scratch[0..n], a);
var best: T = scratch[0];
var best_run: usize = 0;
var run: usize = 0;
for (0..n) |i| {
run = if (i > 0 and scratch[i] == scratch[i - 1]) run + 1 else 1;
// Strictly greater: the sorted order means the first run of a given length starts at the
// smallest value, so ties resolve downward without a second comparison.
if (run > best_run) {
best_run = run;
best = scratch[i];
}
}
return best;
}
valueCounts/// The distinct values of `a` and how often each occurs, most common first.
///
/// Writes into `out_values` and `out_counts` and returns how many distinct values there were, so
/// the caller sizes the buffers for the worst case and reads back the truth. Ties in frequency
/// are broken by value, ascending, for the same reason `mode` does it.
pub fn valueCounts(
comptime T: type,
out_values: []T,
out_counts: []usize,
scratch: []T,
a: Tensor(T),
) Error!usize {
comptime requireFloat(T);
const n: usize = a.size();
if (scratch.len < n or out_values.len < n or out_counts.len < n) {
return Error.ShapeMismatch;
}
try sort(T, scratch[0..n], a);
var distinct: usize = 0;
var i: usize = 0;
while (i < n) {
var run: usize = 1;
while (i + run < n and scratch[i + run] == scratch[i]) {
run += 1;
}
out_values[distinct] = scratch[i];
out_counts[distinct] = run;
distinct += 1;
i += run;
}
// Insertion sort by descending count, ascending value - stable and short, and the list is
// the number of DISTINCT values, which is usually far smaller than the input.
var head: usize = 1;
while (head < distinct) : (head += 1) {
const value: T = out_values[head];
const times: usize = out_counts[head];
var slot: usize = head;
while (slot > 0 and (out_counts[slot - 1] < times or
(out_counts[slot - 1] == times and out_values[slot - 1] > value))) : (slot -= 1)
{
out_values[slot] = out_values[slot - 1];
out_counts[slot] = out_counts[slot - 1];
}
out_values[slot] = value;
out_counts[slot] = times;
}
return distinct;
}
correlationMatrix/// Every pairwise correlation between `columns`, written into a square `out`.
///
/// The diagonal is exactly 1 - not computed, ASSIGNED, because a column's correlation with
/// itself is one by definition and letting the arithmetic produce 0.9999999 would invite a
/// reader to wonder. The matrix is symmetric, so only one triangle is computed and mirrored.
pub fn correlationMatrix(
comptime T: type,
out: Tensor(T),
columns: []const Tensor(T),
) Error!void {
comptime requireFloat(T);
if (out.rank != 2 or out.shape[0] != columns.len or out.shape[1] != columns.len) {
return Error.ShapeMismatch;
}
for (0..columns.len) |i| {
try out.setAt(&.{ i, i }, 1);
for (i + 1..columns.len) |j| {
const r: T = try correlation(T, columns[i], columns[j]);
try out.setAt(&.{ i, j }, r);
try out.setAt(&.{ j, i }, r);
}
}
}
medianAbsDev and meanAbsDev are two functions, not one with a
flag. "MAD" means both, and znum takes a MadKind to disambiguate — so every
call site reads mad(x, .median) and every reader has to know the enum. Two names say
it without one, and the median version is the robust one: replacing a single value
1000 and watches the standard deviation multiply by 282× while the median
deviation does not move at all. That gap is the whole reason the function exists.
Ties in mode go to the smallest value, and that is a decision the
caller should not have to guess. Two values can be equally common; returning whichever the scan
met first makes the answer depend on element order — reproducible until the day it is not. The
smallest is arbitrary too, but it is arbitrary the same way every time. valueCounts
uses the same rule, so the first value it reports is the mode and
exactly that rather than the two agreeing by luck.
coefficientOfVariation and standardError are spelled out where znum
has cv and sem. Two letters standing for three words is a lookup, and
these are read far more often than they are typed.
A mean of zero gives the spread nothing to be a fraction of, so
coefficientOfVariation returns DomainError rather than an infinity —
the question does not apply, which is different from the answer being large.
One number is worth reading. Four times the data should halve the standard error,
and sqrt(96/24) is exactly 2 — but with .sample the measured ratio is
2.032. The gap is Bessel's correction: dividing by n - 1 inflates the
smaller sample proportionally more. Asserting 2.0 to a tight bar would have been asserting that
Bessel does nothing and 2.032 for .sample and exactly 2.0 for
.population, which is the cleaner way to state the law.
correlation takes no ddof. It cancels between the numerator and both
deviations; znum takes one anyway and its own doc admits it "genuinely doesn't matter". One fewer
decision at the call site is one fewer thing to get wrong.
Suppose a column holds repeated values from a small set — material ids, enemy states, weapon types, the winner of each round. You want to replace each one with a small integer, and you want to be able to get the original back.
const levels = try zn.factorize(f64, &codes, &out_levels, column, .sorted);
// codes[i] is a small integer; out_levels[codes[i]] is the original valuefactorize/// Replace each value with a small integer code, and report the distinct values.
///
/// THE CALLER CHOOSES THE ORDER, BECAUSE BOTH ARE RIGHT AND THE WRONG ONE IS SILENT
///
/// znum assigns codes in order of first appearance, which is pandas' default. My first version
/// here sorted instead and the register row claimed sorting was better. **That was too strong**,
/// and understanding why pandas chose first-appearance is what corrected it:
///
/// - **First-appearance needs only equality.** pandas factorizes strings, tuples and objects,
/// many of which have no total order or an expensive locale-dependent one. Sorting needs a
/// comparison. That is a capability difference, not a preference - though it does not bind
/// here, where `T` is always numeric.
/// - **It is the cheaper pass**, and if you are compressing a column or grouping by it, the
/// code VALUES do not matter - only that equal values share one.
/// - **And pandas offers both**: `sort=True` gives the sorted order. It defaults to the cheaper
/// one that works on more types, which is a reasonable default and not a claim that order of
/// appearance is correct.
///
/// The failure I argued from - train on one shuffle, serve with another - needs the encoder to be
/// RE-DERIVED at serving time. A pipeline that persists the levels, which is what this function
/// returns them for, is reproducible under either order. That was the hole in the argument.
///
/// What survives: factorizing a train split and a test split separately is a real and common
/// mistake, and under `.first_seen` the two halves silently disagree about what code 0 means.
/// Under `.sorted` they agree whenever they contain the same values.
///
/// So neither is a default here. Six other places in this file take a named argument for exactly
/// this reason - `Use`, `Look`, `Push`, `FeedForward`, `Comparison`, `Ddof` - and this is the
/// seventh: both answers are defensible, and choosing wrong produces plausible numbers.
///
/// Returns how many distinct values there were. `out_levels[out_codes[i]]` recovers the input
/// under either order, which is the inverse the encoding exists to have.
pub fn factorize(
comptime T: type,
out_codes: []usize,
out_levels: []T,
a: Tensor(T),
order: CodeOrder,
) Error!usize {
comptime requireNumeric(T);
const n: usize = a.size();
if (out_codes.len < n or out_levels.len < n) {
return Error.ShapeMismatch;
}
switch (order) {
.sorted => {
// `unique` sorts into `out_levels` and returns how many are distinct - no scratch
// needed, which is one buffer fewer than the caller would otherwise have to size.
const distinct: usize = try unique(T, out_levels, a);
var walk: Walk = .over(a.shape[0..a.rank]);
var position: usize = 0;
while (walk.next()) |at| : (position += 1) {
const value: T = try a.at(at);
// A binary search, which is what sorted levels buy over a linear scan.
var low: usize = 0;
var high: usize = distinct;
while (low < high) {
const middle: usize = low + (high - low) / 2;
if (out_levels[middle] < value) {
low = middle + 1;
} else {
high = middle;
}
}
out_codes[position] = low;
}
return distinct;
},
.first_seen => {
// A linear scan over the levels seen so far. A categorical column has few levels, so
// this stays effectively linear in `n` - znum makes the same observation, and reaching
// for a hash map would be solving a problem neither of us has.
var distinct: usize = 0;
var walk: Walk = .over(a.shape[0..a.rank]);
var position: usize = 0;
while (walk.next()) |at| : (position += 1) {
const value: T = try a.at(at);
var code: ?usize = null;
for (out_levels[0..distinct], 0..) |level, k| {
if (level == value) {
code = k;
break;
}
}
if (code) |found| {
out_codes[position] = found;
} else {
out_levels[distinct] = value;
out_codes[position] = distinct;
distinct += 1;
}
}
return distinct;
},
}
}
codes gets one integer per input element. out_levels gets the distinct
values, one per code, so out_levels[codes[i]] recovers column[i]. The
return is how many distinct values there were.
That is useful for three ordinary reasons: the codes index an array directly, they are cheap to compare and store, and a neural network can only eat numbers — so a column of "which of six enemy types" becomes a column of 0 to 5.
Forced: equal values must share a code, and the inverse must recover the input. Both hold whatever else you decide.
Chosen: which integer each value gets. Two answers are in common use and neither is wrong:
.first_seen the value appearing first gets 0
.sorted the smallest value gets 0On the same six values in two arrangements:
.sorted 30 10 20 10 30 30 -> 2 0 1 0 2 2 levels 10 20 30
10 30 30 20 10 30 -> 0 2 2 1 0 2 levels 10 20 30
.first_seen 30 10 20 10 30 30 -> 0 1 2 1 0 0 levels 30 10 20
10 30 30 20 10 30 -> 0 1 1 2 0 1 levels 10 30 20Under .sorted the codes depend only on which values are present. Under
.first_seen they depend on the order the rows happen to be in — rearrange the
data and 30 goes from code 0 to code 1.
That matters if you ever number two batches separately and then compare them: under
.first_seen the two can disagree about what 0 means, and nothing will tell you. It does
not matter at all if you number once and keep the levels, which is what they are returned for.
.first_seen is the cheaper pass and the one that would work on values with no
natural order (text, say). .sorted costs a sort and gives an answer that cannot move.
Neither dominates, so the argument is required — this library asks the caller
to choose whenever both answers are defensible and picking the wrong one produces plausible-looking
numbers. Ddof, Use, Look and FeedForward are the
same idea.
This is pd.factorize. pandas defaults to first appearance and offers
sort=True; here there is no default. Its reason for that default is worth knowing:
first appearance needs only equality, so pandas can number strings and arbitrary objects,
where sorting needs a comparison. That does not bind here, where the element type is always a
number — which is why this library asks rather than defaults.
It works at i32 as well as f64 — which is what such a column
usually is, and could not have worked before integer element types were widened, because
unique refused non-floats.
A Series is one named column of numbers, plus an optional record of which
entries are actually there.
const col = try zn.Series(f64).fromSlice(gpa, "price", &.{ 1, 2, 3 });
col.valid == null // nothing missing, so no mask was allocatedSeries/// One named column, with an optional record of which entries are actually there.
///
/// FOLLOWING znum HERE, BECAUSE ITS LAYOUT IS RIGHT
///
/// `valid == null` means **every value is present**, so an ordinary complete column costs no
/// mask and no allocation. Only a column with a real gap pays for one. That is znum's design and
/// it is the thing to copy: the alternative - always allocating a mask - taxes the common case
/// to make the rare one uniform.
///
/// MISSING IS NOT NaN, AND FOR AN INTEGER COLUMN IT CANNOT BE
///
/// A NaN is a float value that arithmetic produces and propagates. Missing is an ABSENCE: nobody
/// measured it, and a reduction should skip it rather than poison itself with it.
///
/// For floats you can almost get away with conflating them. **For integers there is no NaN at
/// all**, and the value a gap would otherwise take is zero - which is a number someone might
/// believe. The integer widening earlier in this port is what made that concrete: `diff`'s first
/// hole is NaN in a float column and 0 in an integer one, and only one of those announces itself.
///
/// So the mask is the representation and NaN is a source you can convert FROM, which is what
/// `fromTensorMarkingNaN` is for.
pub fn Series(comptime T: type) type {
return struct {
const Self = @This();
/// Borrowed - whoever made the column keeps the string alive.
name: []const u8,
/// Rank 1 and contiguous. The column owns it.
values: Tensor(T),
/// `null` means every value is present. Otherwise `valid[i]` says whether entry `i` is
/// real. **Never allocated unless something is actually missing.**
valid: ?[]bool,
/// Who allocated `values.data` and `valid`, so `deinit` can hand both back.
gpa: Allocator,
/// A complete column from raw values: no gaps, no mask, one allocation.
pub fn fromSlice(gpa: Allocator, name: []const u8, data: []const T) Error!Self {
comptime requireNumeric(T);
const values: Tensor(T) = try Tensor(T).alloc(gpa, &.{data.len});
@memcpy(values.data, data);
return .{ .name = name, .values = values, .valid = null, .gpa = gpa };
}
/// A column from a tensor, treating every NaN in it as MISSING.
///
/// The bridge from float data that uses NaN as a gap marker - what `diff` leaves in its
/// first slot, or what a reader produces from an empty CSV field - into an honest mask
/// that reductions skip instead of a NaN that would poison them.
///
/// **Only allocates a mask if something is actually missing**, so a clean column stays
/// in the cheap representation.
pub fn fromTensorMarkingNaN(gpa: Allocator, name: []const u8, source: Tensor(T)) Error!Self {
comptime requireFloat(T);
if (source.rank != 1) {
return Error.UnsupportedShape;
}
const n: usize = source.size();
const values: Tensor(T) = try Tensor(T).alloc(gpa, &.{n});
var gaps: usize = 0;
for (0..n) |i| {
const x: T = try source.at(&.{i});
values.data[i] = x;
if (isNan(x)) {
gaps += 1;
}
}
if (gaps == 0) {
return .{ .name = name, .values = values, .valid = null, .gpa = gpa };
}
const mask: []bool = gpa.alloc(bool, n) catch {
gpa.free(values.data);
return Error.OutOfMemory;
};
for (0..n) |i| {
mask[i] = !isNan(values.data[i]);
}
return .{ .name = name, .values = values, .valid = mask, .gpa = gpa };
}
pub fn deinit(self: Self) void {
self.gpa.free(self.values.data);
if (self.valid) |mask| {
self.gpa.free(mask);
}
}
/// How many entries there are, present or not.
pub fn len(self: Self) usize {
return self.values.size();
}
/// Whether entry `i` is present. A column with no mask says yes to everything, which is
/// what `null` means.
pub fn isValid(self: Self, i: usize) bool {
const mask: []bool = self.valid orelse return true;
return mask[i];
}
/// The sum of the present entries, and how many there were.
///
/// SKIPPING IS THE WHOLE POINT OF THE MASK
///
/// `sumAll` over the raw values would add whatever sits in the gaps - NaN in a float
/// column, which poisons the total, or **zero in an integer one, which does not** and is
/// therefore worse: the answer looks fine and is wrong by however many gaps there were.
///
/// Returns the count too, because a mean needs it and recomputing it means walking the
/// mask twice. A column with nothing present sums to zero over zero, and the caller
/// decides what that means rather than getting a NaN from here.
pub fn sumValid(self: Self) struct { total: T, count: usize } {
var total: CompensatedSum(T) = .{};
var seen: usize = 0;
for (0..self.len()) |i| {
if (self.isValid(i)) {
total.add(self.values.at1(i));
seen += 1;
}
}
return .{ .total = total.value(), .count = seen };
}
/// The mean of the present entries, or null if none are.
///
/// **Null rather than NaN**: an empty mean has no answer, and saying so in the type is
/// what stops the absence propagating silently into whatever comes next.
pub fn meanValid(self: Self) ?T {
const summed = self.sumValid();
if (summed.count == 0) {
return null;
}
return switch (@typeInfo(T)) {
.int => @divTrunc(summed.total, @as(T, @intCast(summed.count))),
else => summed.total / @as(T, @floatFromInt(summed.count)),
};
}
/// Drop the mask if nothing is missing any more.
///
/// After a fill, a column may be complete again - and a complete column should be back
/// in the cheap representation, where `valid == null` means "everything is here" and no
/// reduction pays to consult a mask. znum does this and it is worth copying: **a
/// representation you can only enter and never leave is a one-way ratchet.**
fn dropMaskIfComplete(self: *Self) void {
const mask: []bool = self.valid orelse return;
for (mask) |present| {
if (!present) {
return;
}
}
self.gpa.free(mask);
self.valid = null;
}
/// Fill every gap with a constant, in place. pandas calls this `fillna(value)`.
///
/// The column is complete afterwards by construction, so the mask is freed rather than
/// left behind saying everything is true.
pub fn fillMissing(self: *Self, value: T) void {
const mask: []bool = self.valid orelse return;
for (0..self.len()) |i| {
if (!mask[i]) {
self.values.setAt1(i, value);
}
}
self.gpa.free(mask);
self.valid = null;
}
/// Each gap takes the last real value before it. pandas calls this `ffill`.
///
/// **A LEADING GAP STAYS MISSING**, because there is nothing earlier to borrow. That is
/// not an edge case to paper over: inventing a value for it - zero, or the first real
/// value found later - would be fabricating data at exactly the point where the series
/// has none, and the column would stop being able to say so.
///
/// THE VALUE UNDER A MISSING SLOT IS UNSPECIFIED. A fill writes where it fills and
/// nowhere else, so an unfilled slot keeps whatever it held - the NaN it was built from,
/// or for an integer column whatever the tensor contained. **Only the mask is
/// authoritative**; read `isValid` before reading a value, and every reduction here
/// does.
pub fn fillForward(self: *Self) void {
if (self.valid == null) {
return;
}
var carried: ?T = null;
// `at1`/`setAt1`, not `data[i]` - a `Series` may sit over a STRIDED view, and
// reading the buffer directly walks the wrong elements. Measured: a column taken
// from a matrix summed to 21 where the answer was 7.
for (0..self.len()) |i| {
if (self.valid.?[i]) {
carried = self.values.at1(i);
} else if (carried) |value| {
self.values.setAt1(i, value);
self.valid.?[i] = true;
}
}
self.dropMaskIfComplete();
}
/// Each gap takes the next real value after it. pandas calls this `bfill`.
///
/// The mirror of `fillForward`, and a TRAILING gap stays missing for the same reason.
pub fn fillBackward(self: *Self) void {
if (self.valid == null) {
return;
}
var carried: ?T = null;
var i: usize = self.values.data.len;
while (i > 0) {
i -= 1;
if (self.valid.?[i]) {
carried = self.values.at1(i);
} else if (carried) |value| {
self.values.setAt1(i, value);
self.valid.?[i] = true;
}
}
self.dropMaskIfComplete();
}
/// A NEW column holding only the present entries, in order.
///
/// Not in place, because dropping rows changes the LENGTH - and a column that shrinks
/// under its neighbours is how a `Frame` becomes ragged. The caller gets a separate
/// column and decides what to do with it; `Frame.addColumn` will refuse to put it back
/// beside columns of the original height, which is the intended outcome.
pub fn dropMissing(self: Self, gpa: Allocator, new_name: []const u8) Error!Self {
const kept: usize = self.countValid();
const values: Tensor(T) = try Tensor(T).alloc(gpa, &.{kept});
var at: usize = 0;
for (0..self.len()) |i| {
if (self.isValid(i)) {
values.data[at] = self.values.at1(i);
at += 1;
}
}
return .{ .name = new_name, .values = values, .valid = null, .gpa = gpa };
}
/// How many entries are actually present.
pub fn countValid(self: Self) usize {
const mask: []bool = self.valid orelse return self.len();
var total: usize = 0;
for (mask) |present| {
if (present) {
total += 1;
}
}
return total;
}
};
}
valid == null means every value is present. An ordinary complete
column costs no mask and no extra allocation; only a column with a real gap pays for one. That is
znum's design and it is the right one — always allocating a mask would tax the common case to
make the rare one uniform.
A NaN is a floating-point value. Arithmetic produces it and propagates it. Missing is an absence: nobody measured that entry, and a sum should skip it rather than inherit it.
For floats you can almost get away with conflating them, because a NaN at least makes the total loudly wrong:
Say a sensor logged five readings and two of them failed:
readings 1, NaN, 3, NaN, 5
sumAll -> NaN poisoned, but you can SEE it
sumValid -> 9 over 3 skipped, so the mean is 3 and not 1.8For an integer column there is no NaN at all. Whatever sits in a gap is an ordinary number — zero, if nobody filled it in — and summing it is silent:
Now the same shape with a counter rather than a sensor — daily visits, where the logger was down on two days:
visits 10, 0, 30, 0, 50 the two zeros were never measured
raw mean 90 / 5 = 18 plausible, and wrong
meanValid 90 / 3 = 30 correctAn 18 and a 30 both look like reasonable daily averages. Nothing about the 18 announces itself — and "the logger was down" and "nobody visited" are the two readings of a zero that the mask exists to tell apart.
No NaN, no error, nothing to notice. That is the case that makes the mask a requirement rather than a convenience, and it only became expressible here once integer element types were widened — before that a column like this could not exist.
One smaller decision follows from the same thinking: meanValid returns
null, not NaN, when nothing is present. An empty mean has no answer, and saying so
in the type is what stops the absence travelling onward looking like a number.
var table = zn.Frame.init(gpa);
defer table.deinit();
try table.addColumn(.{ .f64_values = try zn.Series(f64).fromSlice(gpa, "price", &.{...}) });
try table.addColumn(.{ .i32_values = try zn.Series(i32).fromSlice(gpa, "visits", &.{...}) });Frame/// A table: named columns of possibly different types, all the same height.
///
/// A RAGGED TABLE IS A BUG, SO `addColumn` REFUSES ONE
///
/// Every column must be the same height. Accepting a short one and letting the disagreement
/// surface later - as a bounds failure in some reduction three functions away - is the kind of
/// thing that costs an afternoon. znum refuses it at the door and so does this.
pub const Frame = struct {
gpa: Allocator,
columns: std.ArrayList(Column),
pub fn init(gpa: Allocator) Frame {
return .{ .gpa = gpa, .columns = .empty };
}
pub fn deinit(self: *Frame) void {
for (self.columns.items) |held| {
held.deinit();
}
self.columns.deinit(self.gpa);
}
/// How many columns.
pub fn width(self: *const Frame) usize {
return self.columns.items.len;
}
/// How many rows. Every column is the same height, so the first one answers for all.
pub fn height(self: *const Frame) usize {
if (self.columns.items.len == 0) {
return 0;
}
return self.columns.items[0].len();
}
/// Take ownership of a column and append it.
///
/// Refuses a height that disagrees with what is already there, and refuses a duplicate name -
/// **two columns called "price" make every lookup ambiguous**, and a lookup that silently
/// returns the first of two is worse than one that fails.
pub fn addColumn(self: *Frame, taken: Column) Error!void {
if (self.columns.items.len > 0 and taken.len() != self.height()) {
return Error.ShapeMismatch;
}
if (self.find(taken.name()) != null) {
return Error.DomainError;
}
self.columns.append(self.gpa, taken) catch return Error.OutOfMemory;
}
/// The index of the column with this name, or null.
pub fn find(self: *const Frame, wanted: []const u8) ?usize {
for (self.columns.items, 0..) |held, i| {
if (std.mem.eql(u8, held.name(), wanted)) {
return i;
}
}
return null;
}
/// Give a column a new name.
///
/// Refuses a name another column already has, for the reason `addColumn` does: two columns
/// called "price" make every lookup ambiguous, and renaming into a collision is the same
/// mistake arriving by a different door.
pub fn rename(self: *Frame, from: []const u8, to: []const u8) Error!void {
const at: usize = self.find(from) orelse return Error.DomainError;
if (self.find(to)) |taken| {
// Renaming a column to what it already is should not be an error - it is a no-op,
// and refusing it would make an idempotent script fail on its second run.
if (taken != at) {
return Error.DomainError;
}
return;
}
switch (self.columns.items[at]) {
inline else => |*held| held.name = to,
}
}
/// Remove a column and free it.
///
/// The remaining columns keep their order, which matters because a caller who found a column
/// at index 3 before the drop expects index 3 to still mean something after it. Removing from
/// the middle shifts the later ones down by one - stated because the alternative, swapping
/// the last into the hole, is faster and reorders the table silently.
pub fn dropColumn(self: *Frame, name: []const u8) Error!void {
// `if (... == null) return` rather than `orelse return`, and not for style: the
// `useless-error-return` rule reads a bare `return Error.X` STATEMENT as proof the body
// can error, and does not see one buried in an `orelse`. Its body says so - the check is
// an approximation and offers `lint:off` for what it misses.
//
// Restructuring beats suppressing. The plain form is also what `rename` above uses, so
// the two now read alike.
const found: ?usize = self.find(name);
if (found == null) {
return Error.DomainError;
}
const at: usize = found.?;
self.columns.items[at].deinit();
_ = self.columns.orderedRemove(at);
}
/// The column with this name, or null.
pub fn columnNamed(self: *const Frame, wanted: []const u8) ?Column {
const at: usize = self.find(wanted) orelse return null;
return self.columns.items[at];
}
};
Series/// One named column, with an optional record of which entries are actually there.
///
/// FOLLOWING znum HERE, BECAUSE ITS LAYOUT IS RIGHT
///
/// `valid == null` means **every value is present**, so an ordinary complete column costs no
/// mask and no allocation. Only a column with a real gap pays for one. That is znum's design and
/// it is the thing to copy: the alternative - always allocating a mask - taxes the common case
/// to make the rare one uniform.
///
/// MISSING IS NOT NaN, AND FOR AN INTEGER COLUMN IT CANNOT BE
///
/// A NaN is a float value that arithmetic produces and propagates. Missing is an ABSENCE: nobody
/// measured it, and a reduction should skip it rather than poison itself with it.
///
/// For floats you can almost get away with conflating them. **For integers there is no NaN at
/// all**, and the value a gap would otherwise take is zero - which is a number someone might
/// believe. The integer widening earlier in this port is what made that concrete: `diff`'s first
/// hole is NaN in a float column and 0 in an integer one, and only one of those announces itself.
///
/// So the mask is the representation and NaN is a source you can convert FROM, which is what
/// `fromTensorMarkingNaN` is for.
pub fn Series(comptime T: type) type {
return struct {
const Self = @This();
/// Borrowed - whoever made the column keeps the string alive.
name: []const u8,
/// Rank 1 and contiguous. The column owns it.
values: Tensor(T),
/// `null` means every value is present. Otherwise `valid[i]` says whether entry `i` is
/// real. **Never allocated unless something is actually missing.**
valid: ?[]bool,
/// Who allocated `values.data` and `valid`, so `deinit` can hand both back.
gpa: Allocator,
/// A complete column from raw values: no gaps, no mask, one allocation.
pub fn fromSlice(gpa: Allocator, name: []const u8, data: []const T) Error!Self {
comptime requireNumeric(T);
const values: Tensor(T) = try Tensor(T).alloc(gpa, &.{data.len});
@memcpy(values.data, data);
return .{ .name = name, .values = values, .valid = null, .gpa = gpa };
}
/// A column from a tensor, treating every NaN in it as MISSING.
///
/// The bridge from float data that uses NaN as a gap marker - what `diff` leaves in its
/// first slot, or what a reader produces from an empty CSV field - into an honest mask
/// that reductions skip instead of a NaN that would poison them.
///
/// **Only allocates a mask if something is actually missing**, so a clean column stays
/// in the cheap representation.
pub fn fromTensorMarkingNaN(gpa: Allocator, name: []const u8, source: Tensor(T)) Error!Self {
comptime requireFloat(T);
if (source.rank != 1) {
return Error.UnsupportedShape;
}
const n: usize = source.size();
const values: Tensor(T) = try Tensor(T).alloc(gpa, &.{n});
var gaps: usize = 0;
for (0..n) |i| {
const x: T = try source.at(&.{i});
values.data[i] = x;
if (isNan(x)) {
gaps += 1;
}
}
if (gaps == 0) {
return .{ .name = name, .values = values, .valid = null, .gpa = gpa };
}
const mask: []bool = gpa.alloc(bool, n) catch {
gpa.free(values.data);
return Error.OutOfMemory;
};
for (0..n) |i| {
mask[i] = !isNan(values.data[i]);
}
return .{ .name = name, .values = values, .valid = mask, .gpa = gpa };
}
pub fn deinit(self: Self) void {
self.gpa.free(self.values.data);
if (self.valid) |mask| {
self.gpa.free(mask);
}
}
/// How many entries there are, present or not.
pub fn len(self: Self) usize {
return self.values.size();
}
/// Whether entry `i` is present. A column with no mask says yes to everything, which is
/// what `null` means.
pub fn isValid(self: Self, i: usize) bool {
const mask: []bool = self.valid orelse return true;
return mask[i];
}
/// The sum of the present entries, and how many there were.
///
/// SKIPPING IS THE WHOLE POINT OF THE MASK
///
/// `sumAll` over the raw values would add whatever sits in the gaps - NaN in a float
/// column, which poisons the total, or **zero in an integer one, which does not** and is
/// therefore worse: the answer looks fine and is wrong by however many gaps there were.
///
/// Returns the count too, because a mean needs it and recomputing it means walking the
/// mask twice. A column with nothing present sums to zero over zero, and the caller
/// decides what that means rather than getting a NaN from here.
pub fn sumValid(self: Self) struct { total: T, count: usize } {
var total: CompensatedSum(T) = .{};
var seen: usize = 0;
for (0..self.len()) |i| {
if (self.isValid(i)) {
total.add(self.values.at1(i));
seen += 1;
}
}
return .{ .total = total.value(), .count = seen };
}
/// The mean of the present entries, or null if none are.
///
/// **Null rather than NaN**: an empty mean has no answer, and saying so in the type is
/// what stops the absence propagating silently into whatever comes next.
pub fn meanValid(self: Self) ?T {
const summed = self.sumValid();
if (summed.count == 0) {
return null;
}
return switch (@typeInfo(T)) {
.int => @divTrunc(summed.total, @as(T, @intCast(summed.count))),
else => summed.total / @as(T, @floatFromInt(summed.count)),
};
}
/// Drop the mask if nothing is missing any more.
///
/// After a fill, a column may be complete again - and a complete column should be back
/// in the cheap representation, where `valid == null` means "everything is here" and no
/// reduction pays to consult a mask. znum does this and it is worth copying: **a
/// representation you can only enter and never leave is a one-way ratchet.**
fn dropMaskIfComplete(self: *Self) void {
const mask: []bool = self.valid orelse return;
for (mask) |present| {
if (!present) {
return;
}
}
self.gpa.free(mask);
self.valid = null;
}
/// Fill every gap with a constant, in place. pandas calls this `fillna(value)`.
///
/// The column is complete afterwards by construction, so the mask is freed rather than
/// left behind saying everything is true.
pub fn fillMissing(self: *Self, value: T) void {
const mask: []bool = self.valid orelse return;
for (0..self.len()) |i| {
if (!mask[i]) {
self.values.setAt1(i, value);
}
}
self.gpa.free(mask);
self.valid = null;
}
/// Each gap takes the last real value before it. pandas calls this `ffill`.
///
/// **A LEADING GAP STAYS MISSING**, because there is nothing earlier to borrow. That is
/// not an edge case to paper over: inventing a value for it - zero, or the first real
/// value found later - would be fabricating data at exactly the point where the series
/// has none, and the column would stop being able to say so.
///
/// THE VALUE UNDER A MISSING SLOT IS UNSPECIFIED. A fill writes where it fills and
/// nowhere else, so an unfilled slot keeps whatever it held - the NaN it was built from,
/// or for an integer column whatever the tensor contained. **Only the mask is
/// authoritative**; read `isValid` before reading a value, and every reduction here
/// does.
pub fn fillForward(self: *Self) void {
if (self.valid == null) {
return;
}
var carried: ?T = null;
// `at1`/`setAt1`, not `data[i]` - a `Series` may sit over a STRIDED view, and
// reading the buffer directly walks the wrong elements. Measured: a column taken
// from a matrix summed to 21 where the answer was 7.
for (0..self.len()) |i| {
if (self.valid.?[i]) {
carried = self.values.at1(i);
} else if (carried) |value| {
self.values.setAt1(i, value);
self.valid.?[i] = true;
}
}
self.dropMaskIfComplete();
}
/// Each gap takes the next real value after it. pandas calls this `bfill`.
///
/// The mirror of `fillForward`, and a TRAILING gap stays missing for the same reason.
pub fn fillBackward(self: *Self) void {
if (self.valid == null) {
return;
}
var carried: ?T = null;
var i: usize = self.values.data.len;
while (i > 0) {
i -= 1;
if (self.valid.?[i]) {
carried = self.values.at1(i);
} else if (carried) |value| {
self.values.setAt1(i, value);
self.valid.?[i] = true;
}
}
self.dropMaskIfComplete();
}
/// A NEW column holding only the present entries, in order.
///
/// Not in place, because dropping rows changes the LENGTH - and a column that shrinks
/// under its neighbours is how a `Frame` becomes ragged. The caller gets a separate
/// column and decides what to do with it; `Frame.addColumn` will refuse to put it back
/// beside columns of the original height, which is the intended outcome.
pub fn dropMissing(self: Self, gpa: Allocator, new_name: []const u8) Error!Self {
const kept: usize = self.countValid();
const values: Tensor(T) = try Tensor(T).alloc(gpa, &.{kept});
var at: usize = 0;
for (0..self.len()) |i| {
if (self.isValid(i)) {
values.data[at] = self.values.at1(i);
at += 1;
}
}
return .{ .name = new_name, .values = values, .valid = null, .gpa = gpa };
}
/// How many entries are actually present.
pub fn countValid(self: Self) usize {
const mask: []bool = self.valid orelse return self.len();
var total: usize = 0;
for (mask) |present| {
if (present) {
total += 1;
}
}
return total;
}
};
}
A table is named columns, all the same height, each holding one element type. Zig has no dynamic
dispatch, so holding several types in one list means a tagged union —
Column.
The usual cost of a union is a switch per operation with one identical-looking arm
per variant. inline else removes that:
pub fn len(self: Column) usize {
return switch (self) {
inline else => |held| held.len(),
};
}That stamps out one arm per variant at compile time, with held bound to the concrete
Series. It compiles into "if it is an f64 column call Series(f64).len, if
i32 call Series(i32).len" while the source stays one line. znum reaches for the same
construct for the same reason.
A height that disagrees. A ragged table is a bug, and accepting a short column means the disagreement surfaces later as a bounds failure in some reduction three functions away — the kind of thing that costs an afternoon.
A duplicate name. Two columns called "price" make every lookup ambiguous, and a lookup that silently returns the first of two is worse than one that fails.
Both leave the table untouched, so the rejected column is still yours to free.
znum's Column has a text variant for string columns. This one does not:
a string column is not a tensor, its values cannot be summed or scaled, and every numeric operation
would need an arm that returns an error. factorize turns categories into integer codes,
which is a representation the rest of the library can actually use.
A column with gaps sits beside complete ones and only it carries a mask — one gappy column does not tax the rest.
col.fillMissing(0); // a constant into every gap
col.fillForward(); // each gap takes the last real value before it
col.fillBackward(); // ... or the next one after it
const tight = try col.dropMissing(gpa, "name"); // a NEW, shorter columnA sensor that reports every hour but dropped three samples:
_, 2, _, 4, _ fillForward -> _, 2, 2, 4, 4 carry the last reading
fillBackward -> 2, 2, 4, 4, _ or the next oneForward fill is what you want for a sensor: hold the last known value until a new one arrives. Backward fill is for the opposite case — a value that becomes known later and applies retroactively, like a price correction.
The first slot stays missing under a forward fill, and the last stays missing under a backward one. That is the point, not an edge case to paper over. Inventing a value there — zero, or the first real value found later — would be fabricating data at exactly the place the column has none, and it would lose the ability to say so.
Run both and nothing is missing, so the column returns to the cheap representation: the mask is freed rather than left behind saying everything is true. A representation you can enter and never leave is a one-way ratchet.
A fill writes where it fills and nowhere else, so an unfilled slot keeps whatever it held — the NaN it was built from, or for an integer column whatever the tensor contained. Only the mask is authoritative.
Worth stating because it is easy to assume otherwise, and to expect the
leading slot to read 2.0 after a forward fill, and it reads NaN. Read
isValid before reading a value — every reduction here does.
Removing rows changes the length, and a column that shrinks under its neighbours is how
a table becomes ragged. So dropMissing returns a separate column, and
Frame.addColumn will refuse to put it back beside columns of the original height
— which is the intended outcome rather than an inconvenience.
None of this depends on NaN existing. An integer column fills exactly the same way, because the mask is the representation.
Grouping means: given a column of keys, which rows share each value.
const g = try zn.groupBy(i32, gpa, team_column);
for (0..g.groupCount()) |k| {
const rows = g.rows(k); // row indices in group k
}groupBy/// Gather a column's rows by value.
///
/// Built on `factorize(.sorted)` rather than a hash: the codes are already the group numbers, and
/// sorted order means **the groups do not depend on how the rows arrived**. Two runs over the
/// same data in different orders produce the same grouping, which is what makes their outputs
/// comparable.
///
/// A missing key is its own group - `null` is a value people group by, and dropping those rows
/// silently would lose them. The caller gets `null_group` to find it, or no group at all when
/// every key is present.
pub fn groupBy(comptime K: type, gpa: Allocator, key: Series(K)) Error!GroupBy(K) {
comptime requireNumeric(K);
const n: usize = key.len();
const present: usize = key.countValid();
const gaps: usize = n - present;
// FACTORIZE ONLY THE PRESENT VALUES
//
// The value sitting under a missing slot is unspecified - that is the rule `fillForward`
// documents, and it bites here. Handing the raw column to `factorize` counts those leftovers
// as real keys: a column of `{30, _, 30, _, 10}` whose gaps happen to hold 0 comes back with
// THREE distinct keys instead of two, and then a fourth group for the gaps themselves.
//
// **The test caught this**, which is what asserting a group count is for. Compacting first
// costs one pass and removes the whole class.
const compact: Tensor(K) = try Tensor(K).alloc(gpa, &.{present});
defer gpa.free(compact.data);
var at: usize = 0;
for (0..n) |row| {
if (key.isValid(row)) {
compact.data[at] = key.values.at1(row);
at += 1;
}
}
const codes: []usize = gpa.alloc(usize, if (present == 0) 1 else present) catch
return Error.OutOfMemory;
defer gpa.free(codes);
const levels: []K = gpa.alloc(K, if (present == 0) 1 else present) catch
return Error.OutOfMemory;
defer gpa.free(levels);
const distinct: usize = try factorize(K, codes, levels, compact, .sorted);
const groups: usize = distinct + @intFromBool(gaps > 0);
const starts: []usize = gpa.alloc(usize, groups + 1) catch return Error.OutOfMemory;
@memset(starts, 0);
// Count each group, then turn the counts into offsets - the standard two-pass bucket sort,
// which is stable and needs no comparison beyond the one `factorize` already did.
var seen: usize = 0;
for (0..n) |row| {
const g: usize = if (key.isValid(row)) blk: {
defer seen += 1;
break :blk codes[seen];
} else distinct;
starts[g + 1] += 1;
}
for (1..groups + 1) |g| {
starts[g] += starts[g - 1];
}
const order: []usize = gpa.alloc(usize, n) catch return Error.OutOfMemory;
const cursor: []usize = gpa.alloc(usize, groups) catch return Error.OutOfMemory;
defer gpa.free(cursor);
@memcpy(cursor, starts[0..groups]);
var placed: usize = 0;
for (0..n) |row| {
const g: usize = if (key.isValid(row)) blk: {
defer placed += 1;
break :blk codes[placed];
} else distinct;
order[cursor[g]] = row;
cursor[g] += 1;
}
const kept: []K = gpa.alloc(K, groups) catch return Error.OutOfMemory;
@memcpy(kept[0..distinct], levels[0..distinct]);
if (gaps > 0) {
// The missing group has no key value; zero is a placeholder and `rows` is what matters.
kept[distinct] = 0;
}
return .{ .gpa = gpa, .order = order, .starts = starts, .keys = kept };
}
The layout is two arrays. order holds row indices with each group's rows
contiguous; starts says where each group begins. Group g is
order[starts[g]..starts[g+1]].
Say five play sessions, each on some level:
row 0 1 2 3 4
level 3 1 3 2 1
score 900 250 750 400 300Grouping by level gives:
order 1 4 3 0 2 row indices, gathered by level
starts 0 2 3 5
keys 1 2 3
group 0 (level 1) order[0..2] = rows 1, 4 scores 250, 300
group 1 (level 2) order[2..3] = row 3 score 400
group 2 (level 3) order[3..5] = rows 0, 2 scores 900, 750starts has one more entry than there are groups, and the last equals the
row count. That makes groupCount() simply starts.len - 1, and it means the
last group needs no special case.
Without it, every walk over the groups ends in
if (g == last) total else starts[g+1] — the kind of branch that is right nine
times and wrong the tenth. znum does this and it is worth copying exactly.
Groups come out ordered by key, inherited from factorize(.sorted). Two runs over the
same data in different row orders produce the same groups in the same order — which is what
makes their outputs comparable at all. Within a group the rows keep their original order, because
the two-pass bucket sort is stable.
null is a value people group by, and dropping those rows silently would lose them
— the row count would stop adding up and nothing would say why. The gap group sorts last,
after the real keys.
This is where a bug lived. The first version handed the whole column to
factorize, which counted the values sitting under the missing slots as real
keys — and those values are unspecified, as the fill section explains. A column with two real
keys and two gaps came back with four groups instead of three. The test asserted a group count, the
count was wrong, and one compaction pass removed the whole class.
Carrying on with the five sessions from the last section:
const best = try zn.aggregate(i32, .maximum, gpa, grouped, score, "best");
level 1 2 3
best 300 400 900 one row per group, in key orderaggregate/// Reduce every group to one value.
///
/// THE TYPE SURVIVES WHERE IT CAN, WHICH znum'S DOES NOT
///
/// znum's `agg` returns an f64 column for every aggregate. Six of the ten do not need one: a sum
/// of `i32` is an `i32`, and so are the minimum, maximum, first and last. **A count least of
/// all** - it is a count, and past 2^53 an f64 stops being able to hold one exactly, which is the
/// ceiling register row 31 already documents for index tensors.
///
/// `how` is COMPTIME, so asking for a mean of an integer column is a **compile error** naming the
/// problem rather than a silent widening to f64 or a truncating divide. That is the whole benefit
/// of knowing the aggregate at comptime, and it costs the caller nothing: `how` is a literal at
/// every call site anyone would write.
///
/// Missing values are skipped, not counted and not treated as zero - the same rule `sumValid`
/// follows, for the same reason.
pub fn aggregate(
comptime T: type,
comptime how: Aggregate,
gpa: Allocator,
grouped: anytype,
values: Series(T),
name: []const u8,
) Error!Series(T) {
comptime requireNumeric(T);
comptime {
if (how.needsFloat() and @typeInfo(T) != .float) {
@compileError("aggregate: ." ++ @tagName(how) ++
" has no answer in an integer type - use an f64 column, or .sum with your own divide");
}
}
const groups: usize = grouped.groupCount();
const out: Tensor(T) = try Tensor(T).alloc(gpa, &.{groups});
// A group with nothing present has no answer, so it is MISSING rather than zero - which is
// the distinction `Series` exists to keep, and an aggregate is the first place it would be
// tempting to lose it.
var mask: ?[]bool = null;
for (0..groups) |g| {
const rows: []const usize = grouped.rows(g);
var seen: usize = 0;
var total: CompensatedSum(T) = .{};
var low: T = zm.highest(T);
var high: T = zm.lowest(T);
var first_seen: T = 0;
var last_seen: T = 0;
for (rows) |row| {
if (!values.isValid(row)) {
continue;
}
const x: T = values.values.at1(row);
if (seen == 0) {
first_seen = x;
}
last_seen = x;
total.add(x);
low = @min(low, x);
high = @max(high, x);
seen += 1;
}
if (seen == 0) {
if (mask == null) {
mask = gpa.alloc(bool, groups) catch return Error.OutOfMemory;
@memset(mask.?, true);
}
mask.?[g] = false;
out.data[g] = 0;
continue;
}
out.data[g] = switch (how) {
.count => @intCast(seen),
.sum => total.value(),
.minimum => low,
.maximum => high,
.first => first_seen,
.last => last_seen,
.mean => total.value() / @as(T, @floatFromInt(seen)),
};
}
return .{ .name = name, .values = out, .valid = mask, .gpa = gpa };
}
best is a Series(i32), not a Series(f64) — which is
the point of this section.
A sum of i32 is an i32. So are the minimum, maximum, first and last.
And a count least of all — it is a count, and past 253 an
f64 cannot hold one exactly, which is the same ceiling the index-tensor work already ran into.
znum returns an f64 column for every aggregate. Keeping the type means no cast at the call site and no silent widening in the middle.
aggregate: .mean has no answer in an integer type
- use an f64 column, or .sum with your own dividehow is a comptime parameter precisely so the compiler can say this. The alternatives
are a truncating divide (wrong, silently) or a widen to f64 (right, but now every aggregate
returns a float to serve one case). It costs the caller nothing — how is a
literal at every call site anyone would write.
Missing values are skipped, not counted and not treated as zero. If a group has nothing present, the result is missing rather than zero — a zero there is indistinguishable from a real total of zero, and an aggregate is the first place it would be tempting to lose that distinction.
So the aggregate is itself a Series with a mask, and a reduction over it skips the
empty groups in turn. The absence carries.
No variance, stdDev or median. Each needs more than one
pass — a second once the mean is known, or a sort. An enum variant that always
returns an error is worse than an absent one: it looks available, and the caller finds out
at runtime. grouped.rows(g) gives the row indices, so those are three lines at the
call site using functions that already exist.
const pairs = try zn.joinRows(i32, gpa, left_ids, right_ids, .inner);
// each pair is { .left = row, .right = row }, or RowPair.no_row on the unmatched sidejoinRows/// The row pairs of a join, without building the result table.
///
/// WHY THE PAIRING IS SEPARATE FROM THE GATHER
///
/// znum's `join` returns a whole `Frame`: it matches the keys, then materialises every column
/// from both sides. That is the useful end product and it bundles two different jobs.
///
/// The matching is where the semantics live - which rows pair, what `outer` means, where the
/// unmatched go. The gather is a mechanical copy. **Splitting them means the interesting half
/// can be tested directly**, and a caller who wants three columns out of forty does not pay to
/// materialise the other thirty-seven.
///
/// Rows come out in left order, then the right-only rows - which is znum's order and pandas',
/// so a reader can compare an answer against either.
///
/// A MISSING KEY NEVER MATCHES, not even another missing key. Two rows that both failed to
/// record a value are not thereby the same row, and SQL agrees: `NULL = NULL` is not true. They
/// still appear as unmatched rows under `left`, `right` and `outer`.
pub fn joinRows(
comptime K: type,
gpa: Allocator,
left_key: Series(K),
right_key: Series(K),
how: JoinHow,
) Error![]RowPair {
comptime requireNumeric(K);
var pairs: std.ArrayList(RowPair) = .empty;
errdefer pairs.deinit(gpa);
// Group the right side once, so each left row costs a lookup rather than a scan. That turns
// a quadratic join into `n log n` for the grouping plus a binary search per row.
const right_groups: GroupBy(K) = try groupBy(K, gpa, right_key);
defer right_groups.deinit();
const gaps_on_right: bool = right_key.countValid() < right_key.len();
const real_groups: usize = right_groups.groupCount() - @intFromBool(gaps_on_right);
const matched: []bool = gpa.alloc(bool, right_key.len()) catch return Error.OutOfMemory;
defer gpa.free(matched);
@memset(matched, false);
for (0..left_key.len()) |row| {
var found: bool = false;
if (left_key.isValid(row)) {
const wanted: K = left_key.values.at1(row);
// The groups are sorted by key, so this is a binary search.
var low: usize = 0;
var high: usize = real_groups;
while (low < high) {
const middle: usize = low + (high - low) / 2;
if (right_groups.keys[middle] < wanted) {
low = middle + 1;
} else {
high = middle;
}
}
if (low < real_groups and right_groups.keys[low] == wanted) {
for (right_groups.rows(low)) |other| {
pairs.append(gpa, .{ .left = row, .right = other }) catch
return Error.OutOfMemory;
matched[other] = true;
found = true;
}
}
}
if (!found and (how == .left or how == .outer)) {
pairs.append(gpa, .{ .left = row, .right = RowPair.no_row }) catch
return Error.OutOfMemory;
}
}
if (how == .right or how == .outer) {
for (0..right_key.len()) |other| {
if (!matched[other]) {
pairs.append(gpa, .{ .left = RowPair.no_row, .right = other }) catch
return Error.OutOfMemory;
}
}
}
return pairs.toOwnedSlice(gpa) catch return Error.OutOfMemory;
}
znum's join returns a whole Frame: it matches the keys, then
materialises every column from both sides. That is the useful end product, and it bundles two
different jobs.
The matching is where the semantics live — which rows pair, what
outer means, where the unmatched go. The gather is a mechanical copy.
Splitting them means the interesting half can be tested directly, and a caller who wants three
columns out of forty does not pay to materialise the other thirty-seven.
A table of players and a table of their purchases:
players purchases
row id row id item
0 10 0 20 sword
1 20 1 30 shield
2 30 2 30 potion <- player 30 bought twice
3 40 elixir <- a purchase with no player rowinner (1,0) (2,1) (2,2) three rows from three-plus-fourPlayer 20 pairs once. Player 30 pairs twice, because they bought two things. Player 10 bought nothing and purchase 40 has no player, so neither appears in an inner join.
Left row 2 appears twice, because the key 30 is on the right twice. This is the property that surprises people and the reason a join can return more rows than either input had. Rows come out in left order, then the right-only rows — znum's order and pandas'.
Two rows that both failed to record a value are not thereby the same row. SQL agrees:
NULL = NULL is not true.
The trap is that the raw values under two missing slots may well be equal — both
zero, say — and a join that compared values rather than consulting the mask would happily pair
them. The test uses exactly that case: two columns with a missing first entry whose underlying value
is 0 in both. They do not match, and under outer they appear as two separate unmatched
rows.
The pairs feed takeRows, which is the one gather engine: hand it a
list of row indices and it produces a column. Filtering is a shorter list, sorting is a reordered
one, sampling is a random one, and a join's assembly is the pair list — four operations, one
loop. znum makes the same observation, and it is what keeps a dataframe library from growing five
copies of the same code, four of which eventually differ.
const pairs = try zn.joinRows(i32, gpa, left, right, .left);
// ... split the pairs into two index lists ...
const joined_left = try zn.takeRows(i32, gpa, left, left_side);
const joined_right = try zn.takeRows(i32, gpa, right, right_side);joinRows/// The row pairs of a join, without building the result table.
///
/// WHY THE PAIRING IS SEPARATE FROM THE GATHER
///
/// znum's `join` returns a whole `Frame`: it matches the keys, then materialises every column
/// from both sides. That is the useful end product and it bundles two different jobs.
///
/// The matching is where the semantics live - which rows pair, what `outer` means, where the
/// unmatched go. The gather is a mechanical copy. **Splitting them means the interesting half
/// can be tested directly**, and a caller who wants three columns out of forty does not pay to
/// materialise the other thirty-seven.
///
/// Rows come out in left order, then the right-only rows - which is znum's order and pandas',
/// so a reader can compare an answer against either.
///
/// A MISSING KEY NEVER MATCHES, not even another missing key. Two rows that both failed to
/// record a value are not thereby the same row, and SQL agrees: `NULL = NULL` is not true. They
/// still appear as unmatched rows under `left`, `right` and `outer`.
pub fn joinRows(
comptime K: type,
gpa: Allocator,
left_key: Series(K),
right_key: Series(K),
how: JoinHow,
) Error![]RowPair {
comptime requireNumeric(K);
var pairs: std.ArrayList(RowPair) = .empty;
errdefer pairs.deinit(gpa);
// Group the right side once, so each left row costs a lookup rather than a scan. That turns
// a quadratic join into `n log n` for the grouping plus a binary search per row.
const right_groups: GroupBy(K) = try groupBy(K, gpa, right_key);
defer right_groups.deinit();
const gaps_on_right: bool = right_key.countValid() < right_key.len();
const real_groups: usize = right_groups.groupCount() - @intFromBool(gaps_on_right);
const matched: []bool = gpa.alloc(bool, right_key.len()) catch return Error.OutOfMemory;
defer gpa.free(matched);
@memset(matched, false);
for (0..left_key.len()) |row| {
var found: bool = false;
if (left_key.isValid(row)) {
const wanted: K = left_key.values.at1(row);
// The groups are sorted by key, so this is a binary search.
var low: usize = 0;
var high: usize = real_groups;
while (low < high) {
const middle: usize = low + (high - low) / 2;
if (right_groups.keys[middle] < wanted) {
low = middle + 1;
} else {
high = middle;
}
}
if (low < real_groups and right_groups.keys[low] == wanted) {
for (right_groups.rows(low)) |other| {
pairs.append(gpa, .{ .left = row, .right = other }) catch
return Error.OutOfMemory;
matched[other] = true;
found = true;
}
}
}
if (!found and (how == .left or how == .outer)) {
pairs.append(gpa, .{ .left = row, .right = RowPair.no_row }) catch
return Error.OutOfMemory;
}
}
if (how == .right or how == .outer) {
for (0..right_key.len()) |other| {
if (!matched[other]) {
pairs.append(gpa, .{ .left = RowPair.no_row, .right = other }) catch
return Error.OutOfMemory;
}
}
}
return pairs.toOwnedSlice(gpa) catch return Error.OutOfMemory;
}
takeRows/// A new column holding the named rows, in the order given.
///
/// Indices may repeat and may reorder freely - that is what makes one function serve filtering,
/// sorting, sampling and a join's gather. `RowPair.no_row` produces a MISSING entry, which is
/// what an unmatched join row needs.
pub fn takeRows(
comptime T: type,
gpa: Allocator,
column: Series(T),
wanted: []const usize,
) Error!Series(T) {
comptime requireNumeric(T);
const values: Tensor(T) = try Tensor(T).alloc(gpa, &.{wanted.len});
var mask: ?[]bool = null;
errdefer gpa.free(values.data);
for (wanted, 0..) |row, at| {
// A GAP IN THE SOURCE STAYS A GAP, and so does a `no_row`. Two different reasons for an
// absent value - "this row had none" and "there was no row" - and the result cannot tell
// them apart, which is correct: downstream neither is a number.
const absent: bool = row == RowPair.no_row or !column.isValid(row);
if (absent) {
if (mask == null) {
mask = gpa.alloc(bool, wanted.len) catch return Error.OutOfMemory;
@memset(mask.?, true);
}
mask.?[at] = false;
values.data[at] = 0;
} else {
values.data[at] = column.values.at1(row);
}
}
return .{ .name = column.name, .values = values, .valid = mask, .gpa = gpa };
}
no_row in the list produces a missing entry, which is exactly what an
unmatched row needs. So a left join comes out with every left value present and gaps on the right
— and the gaps are real gaps, skipped by every reduction rather than zeros to filter
later.
firstRows is the same engine with the list 0, 1, ... n. It is not
called head: that name is a local variable in three algorithms in this file and a
field in Attention, and pandas' spelling is not worth shadowing three working
functions for.
The right side is grouped once, so each left row costs a binary search rather than a scan
— n log n instead of quadratic. And no_row is a sentinel rather
than ?usize: the pairs are a dense array walked in a tight loop, and an optional would
double its size for a case that is rare in an inner join and absent in most others.
try zn.perClassMetrics(f64, &metrics, confusion_matrix);
// metrics[k] = { precision, recall, f1, support }perClassMetrics/// Fill `out` with one `PerClass` per class, from a confusion matrix.
pub fn perClassMetrics(
comptime T: type,
out: []PerClass(T),
matrix: Tensor(T),
) Error!void {
comptime requireFloat(T);
if (matrix.rank != 2 or matrix.shape[0] != matrix.shape[1]) {
return Error.UnsupportedShape;
}
const classes: usize = matrix.shape[0];
if (out.len != classes) {
return Error.ShapeMismatch;
}
for (0..classes) |k| {
const hit: T = matrix.at2(k, k);
var row: T = 0;
var column: T = 0;
for (0..classes) |j| {
row += matrix.at2(k, j);
column += matrix.at2(j, k);
}
const precision: T = if (column == 0) 0 else hit / column;
const recall: T = if (row == 0) 0 else hit / row;
// The HARMONIC mean, which is why F1 is not just an average: it is dragged down by the
// smaller of the two, so a model cannot buy a good F1 with one good half. An arithmetic
// mean of precision 1 and recall 0 is 0.5; the harmonic mean is 0.
const f1: T = if (precision + recall == 0) 0 else 2 * precision * recall / (precision + recall);
out[k] = .{
.precision = precision,
.recall = recall,
.f1 = f1,
.support = @intFromFloat(@round(row)),
};
}
}
With the matrix oriented true-down and predicted-across:
precision = M[k][k] / COLUMN k of everything called k, how much was k
recall = M[k][k] / ROW k of everything that WAS k, how much was foundTen instances, five of each class. A classifier that over-predicts class 1 — calling it for eight of the ten:
recall precision
class 1 1.0 0.625 found every 1, but 3 of its 8 were not
class 0 0.4 1.0 only says 0 when sure, so misses 3A classifier that predicts k for everything has perfect recall on
k and terrible precision. One that predicts k only when certain has the
reverse. Reporting one without the other is how a model looks good.
Precision is a column sum and recall a row sum. Swapping them gives plausible numbers on a symmetric matrix and wrong ones on a lopsided one.
precision 1.0, recall 0.4 arithmetic mean 0.70 F1 0.571
precision 1.0, recall 0.0 arithmetic mean 0.50 F1 0.000It is dragged toward the smaller of the two, so a model cannot buy a good F1 with one good half. The second row is the case that makes it obvious: a classifier that finds none of a class scores zero, not one half.
A class the model never predicted has no precision — the question has no instances to ask about. This reports 0, as scikit-learn does, because NaN would poison a macro average and force every caller to filter.
But 0 from no data is not 0 from bad predictions, and support — how many
instances truly belonged to the class — is what distinguishes them. A class with support 0
and F1 0 was never tested, and a macro average that includes it is averaging an
absence.
histogramEdges returns count + 1 values. That off-by-one is what makes
a plot's last bar vanish or its axis stop one bin short.
The last edge is set to the bound rather than computed, for the same reason
linspace does it: an edge a few ulps below its bound puts a value in no bin at
all.
const yesterday = try zn.shiftRows(f64, gpa, visits, 1);
const smoothed = try zn.rollingMean(f64, gpa, visits, 3);shiftRows/// A new column with every value moved `by` rows later, or earlier if negative.
///
/// THE VACATED ROWS ARE MISSING, WHICH IS WHAT THE MASK WAS FOR
///
/// Shifting down by one means row 0 has nothing to take from. That is not zero and not the old
/// row 0 - it is **absent**, and a `Series` can say so. Filling it with zero would put a number
/// where no measurement exists, which is the mistake the whole validity mask exists to prevent.
///
/// The usual reason to shift is a difference against the past: `y - shiftRows(y, 1)` is
/// yesterday's change, and its first row genuinely has no answer. A zero there would read as
/// "no change" rather than "no data".
///
/// Named `shiftRows`, not `shift`: `shiftLeft` and `shiftRight` are the bitwise operations on
/// integer tensors, and one `shift` beside them would be read as a third.
pub fn shiftRows(
comptime T: type,
gpa: Allocator,
column: Series(T),
by: isize,
) Error!Series(T) {
comptime requireNumeric(T);
const n: usize = column.len();
const values: Tensor(T) = try Tensor(T).alloc(gpa, &.{n});
var mask: ?[]bool = null;
errdefer gpa.free(values.data);
for (0..n) |row| {
const from: isize = @as(isize, @intCast(row)) - by;
const inside: bool = from >= 0 and from < @as(isize, @intCast(n));
const present: bool = inside and column.isValid(@intCast(from));
if (present) {
values.data[row] = column.values.at1(@intCast(from));
} else {
if (mask == null) {
mask = gpa.alloc(bool, n) catch return Error.OutOfMemory;
@memset(mask.?, true);
}
mask.?[row] = false;
values.data[row] = 0;
}
}
return .{ .name = column.name, .values = values, .valid = mask, .gpa = gpa };
}
rollingMean/// The mean of each trailing window of `window` rows.
///
/// THE FIRST `window - 1` ROWS HAVE NO FULL WINDOW, AND ARE MISSING
///
/// A three-row average cannot be computed from one row. pandas leaves those slots NaN by default
/// and so does this - as **missing**, which an integer column can also express.
///
/// The alternative, averaging whatever is available, produces a value that looks like the others
/// and is computed from fewer samples. On a noisy series the first few points then look
/// artificially smooth, which is exactly the shape a rolling mean is being used to judge.
///
/// A missing value inside a window makes that window's result missing too. Skipping it would
/// silently change the denominator, and a mean over four of five samples is not the five-sample
/// mean the caller asked for.
pub fn rollingMean(
comptime T: type,
gpa: Allocator,
column: Series(T),
window: usize,
) Error!Series(T) {
comptime requireFloat(T);
if (window == 0) {
return Error.DomainError;
}
const n: usize = column.len();
const values: Tensor(T) = try Tensor(T).alloc(gpa, &.{n});
var mask: ?[]bool = null;
errdefer gpa.free(values.data);
for (0..n) |row| {
var complete: bool = row + 1 >= window;
var total: CompensatedSum(T) = .{};
if (complete) {
const first: usize = row + 1 - window;
for (first..row + 1) |k| {
if (!column.isValid(k)) {
complete = false;
break;
}
total.add(column.values.at1(k));
}
}
if (complete) {
values.data[row] = total.value() / @as(T, @floatFromInt(window));
} else {
if (mask == null) {
mask = gpa.alloc(bool, n) catch return Error.OutOfMemory;
@memset(mask.?, true);
}
mask.?[row] = false;
values.data[row] = 0;
}
}
return .{ .name = column.name, .values = values, .valid = mask, .gpa = gpa };
}
Shifting down by one means row 0 has nothing to take from. That is not zero and not the old row
0 — it is absent, and a Series can say so.
Here is why it matters. The usual reason to shift is a difference against the past:
visits 10 20 30 40 50
yesterday _ 10 20 30 40
change ? 10 10 10 10The first row genuinely has no answer. A zero there would read as "no change" rather than "no data" — and on a chart those look identical.
Shifting up is the mirror: the last row is the absent one. Shifting past the end leaves nothing present, and that is not an error — "shift by more than you have" is an ordinary thing for a loop bound to produce.
A three-row average cannot be computed from one row. pandas leaves those slots NaN and so does this — as missing, which an integer column can also express.
The tempting alternative is to average whatever is available. That gives a value which looks like the others and is computed from fewer samples, so the first few points come out artificially smooth — exactly the shape a rolling mean is being used to judge.
visits 10 _ 30 40 50 window 3
smoothed _ _ _ _ 40 only the last window is cleanSkipping the gap would silently change the denominator, and a mean over two of three samples is not the three-sample mean the caller asked for.
rename refuses a collision — renaming into an existing name is the same
ambiguity addColumn refuses, arriving by another door. But renaming a column to
its own name is a no-op, not an error: refusing it would make an
idempotent script fail on its second run.
dropColumn preserves the order of what remains. Swapping the last column into the
hole would be faster and would reorder the table silently — so a caller who found a
column at index 3 before the drop still gets something sensible after it.
Kernels are Zig, compiled to SPIR-V and translated to WGSL. A complete elementwise one:
const k = @import("kompute");
pub const config = k.Config{ .max = 1 << 14, .workgroup = 64 };
pub const Buffers = extern struct { // last field is the output
x: [config.max]f32,
out: [config.max]f32,
};
pub const Params = extern struct { // must be a multiple of 16 bytes
count: u32,
_p0: u32 = 0, _p1: u32 = 0, _p2: u32 = 0,
};
pub const g = k.Globals(@This());
const bx = g.bind(.x);
const bout = g.bind(.out);
const params = g.uniform();
pub fn relu(id: u32) void {
if (id >= params.count) return;
bout[id] = @max(bx[id], 0);
}
pub const kernels = [_][:0]const u8{ "relu" };
comptime {
for (kernels) |name| k.installKernelLean(@This(), name);
}Uniform blocks are std140-aligned, so Params must be a multiple of 16 bytes; the
host refuses a pipeline otherwise. Kernels that use a barrier must read their parameters through
g.uniform() rather than a copied context, or the loop bounds derived from them are
not provably uniform and the barrier is rejected.
| Situation | CPU | GPU |
|---|---|---|
| A strided or permuted operand | Free — at follows the strides | Must be materialised; a kernel indexes from the start of a buffer |
| A broadcast operand | A stride-0 view | Strides passed in the uniform, or materialised |
A transposed matmul operand | A stride swap | A separate kernel with the transpose in its index arithmetic |
| A domain error | Error.DomainError | No mechanism; the host must not dispatch the case |
sqrt, log below zero | NaN | Undefined in WGSL |
Every kernel is compared against its zn counterpart on a device, element by
element, over 4096 f32 values. Representative worst absolute differences, measured on
an Adreno phone:
| Operation | Worst difference | In ULP of the row's peak |
|---|---|---|
add, sub, mul, abs, neg, masks, floor, ceil, sign, clamp, lerp, transpose, max | 0 | exact |
sigmoid, tanh, log, sqrt | 1.19e-7 | 1 |
gelu, sin, cos | 1.9e-7 – 2.4e-7 | 1–2 |
matmul (64-term contraction) | 5.7e-6 | < 1 |
atan2 | 3.6e-6 | 9.6 |
sum all (one thread) | 3.8e-5 | 3.7 |
sum all (workgroup tree) | 7.6e-6 | 0.7 |
Exact agreement occurs where the operation is a selection or a single flop. The 1-ULP figures
are the two backends' implementations of a transcendental rounding differently.
atan2 is larger because zm.atan2 is std.math.atan2 on the
host and a polynomial approximation on the GPU. The two summation figures differ because a
workgroup tree is pairwise summation, whose error grows as O(log n) against the single
thread's O(n).
zn.
Every GPU kernel here has a CPU twin, and the sweep runs both on the same input and compares. Until now the twin was a transcription: the same arithmetic retyped in the kernel file.
A transcription is a second definition, and second definitions drift. Two bugs came from exactly that — a kernel writing past its output buffer, and another dividing by the wrong width — and neither was a logic error. Both were retyping.
const state: zn.CartpoleState(f32) = .{ .cart = ba[base], ... };
const stepped = zn.cartpoleStep(f32, state, ...);CartpoleState/// The four numbers a cartpole is. **Named, not positions in a `[4]` tensor.**
///
/// znum indexes `state_data[base + 2]` for the pole angle. Nothing at the call site says which of
/// four slots holds what, and swapping two of them produces a plausible simulation that is wrong.
pub fn CartpoleState(comptime T: type) type {
return struct {
/// Metres from the centre of the track.
cart: T,
/// Metres per second.
cart_rate: T,
/// **Radians** from upright - see the note on `cartpoleStep` for why this one angle is
/// not in turns like the rest of the codebase.
pole_rad: T,
/// Radians per second.
pole_rate_rad: T,
};
}
cartpoleStep/// The cart-and-pole control benchmark: push a cart left or right to keep a pole upright.
///
/// A PURE FUNCTION - THE STATE GOES IN AND THE NEXT STATE COMES OUT
///
/// znum's step takes `state: *Tensor(f32)` and updates it IN PLACE, with a doc comment warning
/// that the tensor "holds the observation for step t only until this returns, so a rollout copies
/// it out BEFORE calling". **That is a footgun stated in prose**: the caller has to remember, and
/// forgetting produces a rollout whose observations are all one step late, which trains and
/// converges to something slightly wrong.
///
/// Returning the next state makes it unforgettable. It also makes the step replayable - the same
/// state and action give the same answer, always - which is what a test needs and what a
/// debugger needs.
///
/// RESET IS A SEPARATE FUNCTION, AND THAT REMOVES TWO FLAGS
///
/// znum folds the reset into the step, so the step needs a `force_reset` parameter, and then a
/// `no_transition` special case because a forced reset produces no reward and no termination.
/// Both disappear when starting an episode is its own call: `cartpoleReset` gives a state, and
/// `cartpoleStep` only ever steps.
///
/// RADIANS HERE, AND THE SUFFIX SAYS SO - BECAUSE THIS ANGLE IS DIFFERENTIATED
///
/// The rest of this codebase moved to turns, and for geometry that was right: a quarter turn is
/// exactly 0.25, and `sinTurns` is exact where `sin(tau/4)` is not.
///
/// **A differentiated angle is the other case.** `d/dtheta sin(theta)` is `cos(theta)` only in
/// radians; in turns it is `tau * cos`, and every derivative in a dynamics system picks up that
/// factor. The cartpole equations are differential equations, so radians are the unit in which
/// they have no constants in them at all.
///
/// The first version of this held the pole in turns and converted the angular acceleration from
/// radians on the way in. That conversion was a place to be wrong, and I was: **the pole fell tau
/// times too fast** and every sign still looked right. Choosing the unit the equations are
/// written in removes the conversion rather than fixing it.
///
/// So: **turns for angles that are drawn, radians for angles that are differentiated**, and the
/// name carries the unit either way. `pole_rad` cannot be mistaken for `pole_turns`.
///
/// The limit is still not a rounded constant. znum writes `0.2095`, which is 12.0035 degrees;
/// `radFromDeg(12)` is twelve degrees exactly as far as the type can express it.
///
/// The dynamics are the standard ones - Barto, Sutton and Anderson 1983 - and the integration is
/// SEMI-IMPLICIT Euler: the position update uses the NEW velocity. Gym's default is explicit
/// Euler, which drifts; znum uses semi-implicit too and is right to.
pub fn cartpoleStep(
comptime T: type,
state: CartpoleState(T),
push: Push,
) CartpoleStep(T) {
comptime requireFloat(T);
const gravity: T = 9.8;
const cart_mass: T = 1.0;
const pole_mass: T = 0.1;
const total_mass: T = cart_mass + pole_mass;
const half_pole: T = 0.5;
const pole_moment: T = pole_mass * half_pole;
const push_force: T = 10.0;
const seconds: T = 0.02;
const force: T = switch (push) {
.left => -push_force,
.right => push_force,
};
// `zm.` because zimrnum has its own `sinRad` and it is a tensor op - the disambiguation the
// two `lint:off` lines at the top of this file exist for.
const lean: T = zm.sinRad(state.pole_rad);
const upright: T = zm.cosRad(state.pole_rad);
// NO CONVERSION ANYWHERE, WHICH IS THE POINT OF CHOOSING RADIANS
//
// These are the textbook equations transcribed, and every quantity is already in the unit the
// derivation uses. An earlier version held the pole in turns and converted here; that line
// was a place to be wrong and it was.
const shared: T = (force + pole_moment * state.pole_rate_rad * state.pole_rate_rad * lean) /
total_mass;
const pole_accel: T = (gravity * lean - upright * shared) /
(half_pole * (4.0 / 3.0 - pole_mass * upright * upright / total_mass));
const cart_accel: T = shared - pole_moment * pole_accel * upright / total_mass;
// Semi-implicit: the position uses the UPDATED velocity, which conserves energy far better
// than the explicit form over a long episode.
const cart_rate: T = state.cart_rate + seconds * cart_accel;
const pole_rate_rad: T = state.pole_rate_rad + seconds * pole_accel;
const next: CartpoleState(T) = .{
.cart = state.cart + seconds * cart_rate,
.cart_rate = cart_rate,
.pole_rad = state.pole_rad + seconds * pole_rate_rad,
.pole_rate_rad = pole_rate_rad,
};
// Twelve degrees, converted rather than rounded - and 2.4 metres of track either side.
const fall_rad: T = zm.radFromDeg(@as(T, 12.0));
const track: T = 2.4;
const failed: bool = @abs(next.cart) > track or @abs(next.pole_rad) > fall_rad;
return .{
.state = next,
// A reward of one for every step survived, which is the standard scoring and makes the
// return equal to the episode length - so "did it learn" reads directly off the number.
.reward = 1.0,
.failed = failed,
};
}
That is the kernel. It imports zimrnum and calls the same function the CPU calls, because zimrnum compiles to SPIR-V. The sweep row now compares a function against itself.
The opt-in is .wants_zimrnum on the kernel's build entry. It is per-kernel because
Zig rejects a --dep for a module the file does not import, so handing zn
to every kernel fails on the ones with no use for it.
One detail worth knowing if you touch the build: --dep attaches to the next
-M. Declaring -Mzn= without a matching --dep zn
before -Mroot= names the module without making it visible to the kernel, and the error
— module "zn" declared but not used — reads like the file failed to import
it when the root module's dependency list was the short one.
try zn.gradient(f64, out, values, h); // central inside, one-sided at the ends
const area = try zn.trapezoid(f64, values, h);
try zn.interp(f64, out, at, xs, ys); // clamps outside the rangegradient/// The central-difference derivative of a sampled function, with one-sided differences at the
/// ends. `spacing` is the distance between samples.
///
/// CENTRAL IN THE INTERIOR, ONE-SIDED AT THE ENDS, which is the only choice that gives a
/// result the same length as the input. The interior error is O(h^2) and the two end points are
/// O(h) - stated here because a caller differentiating twice will see the ends degrade first.
pub fn gradient(
comptime T: type,
out: []T,
values: []const T,
spacing: T,
) Error!void {
comptime requireFloat(T);
if (out.len != values.len) {
return Error.ShapeMismatch;
}
if (values.len < 2 or spacing == 0) {
return Error.DomainError;
}
const n: usize = values.len;
out[0] = (values[1] - values[0]) / spacing;
out[n - 1] = (values[n - 1] - values[n - 2]) / spacing;
var i: usize = 1;
while (i + 1 < n) : (i += 1) {
out[i] = (values[i + 1] - values[i - 1]) / (2 * spacing);
}
}
trapezoid/// The trapezoidal integral of evenly spaced samples.
///
/// Exact for any straight line, by construction, and that is what the test asserts - a rule
/// that is not exact on its own basis functions is not the rule it claims to be.
pub fn trapezoid(comptime T: type, values: []const T, spacing: T) Error!T {
comptime requireFloat(T);
if (values.len < 2) {
return Error.DomainError;
}
var total: CompensatedSum(T) = .{};
var i: usize = 1;
while (i < values.len) : (i += 1) {
total.add((values[i - 1] + values[i]) * 0.5 * spacing);
}
return total.value();
}
interp/// Linear interpolation of `(xs, ys)` at each point of `at`, into `out`. `xs` must be strictly
/// increasing; points outside the range are clamped to the nearest endpoint.
///
/// Clamping rather than extrapolating: a linear extrapolation far outside the data is a number
/// with no support, and returning the endpoint is at least a value the data contains. A caller
/// who wants extrapolation can detect the clamp by comparing `at` with `xs`.
pub fn interp(
comptime T: type,
out: []T,
at: []const T,
xs: []const T,
ys: []const T,
) Error!void {
comptime requireFloat(T);
if (xs.len != ys.len or out.len != at.len) {
return Error.ShapeMismatch;
}
if (xs.len < 2) {
return Error.DomainError;
}
var k: usize = 1;
while (k < xs.len) : (k += 1) {
if (xs[k] <= xs[k - 1]) {
return Error.DomainError;
}
}
for (out, at) |*slot, q| {
if (q <= xs[0]) {
slot.* = ys[0];
continue;
}
if (q >= xs[xs.len - 1]) {
slot.* = ys[ys.len - 1];
continue;
}
// Linear scan: `at` is usually short and usually sorted, and a binary search here would
// be the same speed on the sizes this is used at.
var i: usize = 1;
while (xs[i] < q) : (i += 1) {}
const t: T = (q - xs[i - 1]) / (xs[i] - xs[i - 1]);
slot.* = ys[i - 1] + t * (ys[i] - ys[i - 1]);
}
}
Each rule is exact on the basis it is built from, and that basis is what fixes each one: the
trapezoidal rule on any straight line, the central difference on any quadratic, linear
interpolation on any straight line. A rule that is not exact on its own basis is not the rule it
claims to be, and a test against a general function could not tell. The two end points of
gradient are one-sided and therefore O(h) where the interior is O(h²); the
test asserts they are close but not exact, so a reader knows which points to trust.
const C = zn.ComplexNumber(f64);
try zn.fft(f64, data, -1); // forward, in place, power-of-two length
try zn.ifft(f64, data); // inverse, including the 1/n
try zn.dft(f64, out, in, -1); // the definition, any length
try zn.fftFreq(f64, freqs, spacing);ComplexNumber/// A complex number, as a pair.
///
/// `ComplexNumber` rather than `Complex`, which is a reserved name - `zm` has its own, of a different
/// shape. `extern` so an array of them is exactly interleaved re/im, which
/// is the layout every FFT library and every GPU buffer expects.
pub fn ComplexNumber(comptime T: type) type {
comptime requireFloat(T);
return extern struct {
const Self = @This();
re: T = 0,
im: T = 0,
pub fn add(a: Self, b: Self) Self {
return .{ .re = a.re + b.re, .im = a.im + b.im };
}
pub fn sub(a: Self, b: Self) Self {
return .{ .re = a.re - b.re, .im = a.im - b.im };
}
pub fn mul(a: Self, b: Self) Self {
return .{
.re = a.re * b.re - a.im * b.im,
.im = a.re * b.im + a.im * b.re,
};
}
pub fn magnitude(a: Self) T {
return @sqrt(a.re * a.re + a.im * a.im);
}
};
}
fft/// The radix-2 Cooley-Tukey transform, in place. `data.len` must be a power of two.
///
/// O(n log n), WHICH IS THE WHOLE REASON THE NAME EXISTS
///
/// A DFT of 4096 points is 16.8 million complex multiplies; this is 49 152 - 341 times fewer by
/// operation count, and **630 times faster measured**: 10 transforms of 4096 points take
/// **2 520 ms** through `dft` and **4 ms** through `fft`. The measured ratio beats the operation
/// count because the DFT calls `sin` and `cos` per element pair while the butterflies compute one
/// twiddle per pair of blocks.
///
/// znum's `fft` is the O(n^2) definition, so this is a deliberate divergence, justified by that
/// ratio and by `dft` remaining available under its own name for anyone who wants the reference.
///
/// THE FAST TRANSFORM IS CHECKED AGAINST THE SLOW ONE, element by element, at every power of
/// two the test can afford. That is the only check that does not compare an FFT against another
/// FFT written by the same hand.
///
/// Decimation in time: the input is first permuted into bit-reversed order, after which the
/// butterflies combine adjacent pairs, then pairs of pairs, and so on. The permutation is what
/// lets every stage read and write the same array.
pub fn fft(comptime T: type, data: []ComplexNumber(T), exponent_sign: T) Error!void {
comptime requireFloat(T);
const n: usize = data.len;
if (n == 0 or (n & (n - 1)) != 0) {
return Error.UnsupportedShape;
}
if (n == 1) {
return;
}
// Bit-reversal permutation.
var i: usize = 1;
var j: usize = 0;
while (i < n) : (i += 1) {
var bit: usize = n >> 1;
while (j & bit != 0) : (bit >>= 1) {
j ^= bit;
}
j |= bit;
if (i < j) {
const tmp: ComplexNumber(T) = data[i];
data[i] = data[j];
data[j] = tmp;
}
}
// Butterflies, stage by stage.
var width: usize = 2;
while (width <= n) : (width <<= 1) {
const step: T = exponent_sign * @as(T, tau) / @as(T, @floatFromInt(width));
var block: usize = 0;
while (block < n) : (block += width) {
var k: usize = 0;
while (k < width / 2) : (k += 1) {
const angle: T = step * @as(T, @floatFromInt(k));
const twiddle: ComplexNumber(T) = .{ .re = scalarCos(angle), .im = scalarSin(angle) };
const even: ComplexNumber(T) = data[block + k];
const odd: ComplexNumber(T) = twiddle.mul(data[block + k + width / 2]);
data[block + k] = even.add(odd);
data[block + k + width / 2] = even.sub(odd);
}
}
}
}
ifft/// The inverse transform: `fft` with the opposite sign, divided by `n`.
///
/// The 1/n lives here rather than being split as 1/sqrtn between the two directions. Splitting
/// makes the pair unitary, which matters for a spectrogram and not for a round trip; putting it
/// all on the inverse means the forward transform's bin 0 is the plain sum of the input, which is
/// the convention every reader expects.
pub fn ifft(comptime T: type, data: []ComplexNumber(T)) Error!void {
comptime requireFloat(T);
try fft(T, data, 1);
const scale_by: T = 1 / @as(T, @floatFromInt(data.len));
for (data) |*z| {
z.re *= scale_by;
z.im *= scale_by;
}
}
dft/// The discrete Fourier transform by its definition, in place. `sign` is -1 forward, +1 inverse.
///
/// O(n^2), AND IT EXISTS TO CHECK THE FAST ONE
///
/// This is the definition written out: every output bin is a sum over every input sample. It is
/// unusably slow past a few thousand points and it is the only thing `fft` can be checked
/// against that is not another FFT. znum's `fft` is this function; zimrnum keeps it under its
/// real name and adds the fast one beside it.
pub fn dft(
comptime T: type,
out: []ComplexNumber(T),
input: []const ComplexNumber(T),
exponent_sign: T,
) Error!void {
comptime requireFloat(T);
if (out.len != input.len) {
return Error.ShapeMismatch;
}
const n: usize = input.len;
if (n == 0) {
return Error.DomainError;
}
for (out, 0..) |*slot, bin| {
var re: CompensatedSum(T) = .{};
var im: CompensatedSum(T) = .{};
for (input, 0..) |x, sample| {
const angle: T = exponent_sign * @as(T, tau) * @as(T, @floatFromInt(bin * sample % n)) /
@as(T, @floatFromInt(n));
const c: T = scalarCos(angle);
const s: T = scalarSin(angle);
re.add(x.re * c - x.im * s);
im.add(x.re * s + x.im * c);
}
slot.* = .{ .re = re.value(), .im = im.value() };
}
}
fftFreq/// The frequency, in cycles per sample-spacing, of each bin of an `n`-point transform. Bins past
/// the halfway point are negative frequencies, as they are everywhere else.
pub fn fftFreq(comptime T: type, out: []T, spacing: T) Error!void {
comptime requireFloat(T);
const n: usize = out.len;
if (n == 0 or spacing == 0) {
return Error.DomainError;
}
const denom: T = @as(T, @floatFromInt(n)) * spacing;
for (out, 0..) |*slot, bin| {
const signed: isize = if (bin * 2 < n)
@intCast(bin)
else
@as(isize, @intCast(bin)) - @as(isize, @intCast(n));
slot.* = @as(T, @floatFromInt(signed)) / denom;
}
}
dft is the definition and fft is the fast transform, and both
exist on purpose. dft is the formula transcribed — every output bin a
sum over every input sample — unusable past a few thousand points and the only thing
fft can be checked against that is not another FFT. The test compares them element by
element at every power of two up to 256.
dft, 4 ms through
fft — 630×. The operation counts predict 341×; the rest is
the DFT calling sin and cos per element pair where the butterflies
compute one twiddle per pair of blocks.
ComplexNumber(T) is extern, so an array of them is exactly interleaved
re/im — the layout every FFT library and every GPU buffer expects. The whole
1/n lives on the inverse rather than being split as 1/√n between
the two directions: splitting makes the pair unitary, which matters for a spectrogram and not for
a round trip, and putting it all on the inverse means a forward transform's bin 0 is the plain
sum of the input.
Beyond the round trip, this is exact for that a pure tone puts all its energy in bins
k and n − k and nowhere else. A transform with a sign error, a
wrong twiddle or a botched permutation still returns numbers; it does not return a spectrum with
two spikes in the right places.
try zn.rfft(f64, out, signal, scratch); // out is n/2 + 1, scratch is n/2rfft/// The transform of a real signal, into `out`, which must hold `n/2 + 1` bins.
///
/// HALF THE WORK, BECAUSE A REAL SIGNAL'S SPECTRUM IS SYMMETRIC
///
/// `X[n - k]` is the conjugate of `X[k]` when the input is real, so bins past the halfway point
/// carry no information and computing them is waste. The saving is not the discarded half of the
/// output - it is the transform itself: the `n` real samples are packed into `n/2` complex values
/// (`z[k] = x[2k] + i*x[2k+1]`), ONE transform of half the length is run, and the result is
/// untangled into the spectrum. **Half the transform, not a full transform with half thrown
/// away**, which is what znum's `rfft` does.
///
/// THE UNTANGLING IS THE PART THAT CAN BE WRONG, so it is checked against a full `fft` of the
/// same signal, bin for bin, at every power of two the test can afford. A packed transform that
/// got the twiddle or the conjugate wrong still returns a plausible spectrum.
///
/// `input.len` must be even and `n/2` a power of two.
pub fn rfft(
comptime T: type,
out: []ComplexNumber(T),
input: []const T,
scratch: []ComplexNumber(T),
) Error!void {
comptime requireFloat(T);
const n: usize = input.len;
const half: usize = n / 2;
if (n < 2 or n % 2 != 0 or (half & (half - 1)) != 0) {
return Error.UnsupportedShape;
}
if (out.len != half + 1 or scratch.len != half) {
return Error.ShapeMismatch;
}
// Pack: even samples into the real parts, odd into the imaginary.
for (scratch, 0..) |*z, k| {
z.* = .{ .re = input[2 * k], .im = input[2 * k + 1] };
}
try fft(T, scratch, -1);
// Untangle. `Z[k]` and the conjugate of `Z[half - k]` separate the transforms of the even and
// odd sub-sequences; one twiddle recombines them into `X[k]`.
const step: T = -@as(T, tau) / @as(T, @floatFromInt(n));
for (out, 0..) |*slot, k| {
const a: ComplexNumber(T) = scratch[k % half];
const b: ComplexNumber(T) = scratch[(half - k) % half];
// Even part: (Z[k] + conj(Z[half-k])) / 2. Odd part: (Z[k] - conj(Z[half-k])) / 2i,
// written as a multiply by -i/2 to avoid a complex division.
const even: ComplexNumber(T) = .{ .re = (a.re + b.re) * 0.5, .im = (a.im - b.im) * 0.5 };
const gap: ComplexNumber(T) = .{ .re = (a.re - b.re) * 0.5, .im = (a.im + b.im) * 0.5 };
const odd: ComplexNumber(T) = .{ .re = gap.im, .im = -gap.re };
const angle: T = step * @as(T, @floatFromInt(k));
const twiddle: ComplexNumber(T) = .{ .re = scalarCos(angle), .im = scalarSin(angle) };
slot.* = even.add(twiddle.mul(odd));
}
}
X[n − k] is the conjugate of X[k] when the input is real, so
bins past the halfway point carry no information. The saving is not the discarded half of the
output — it is the transform itself: the n real samples are packed into
n/2 complex values, one transform of half the length is run, and the result
is untangled. Half the transform, not a full transform with half thrown away,
which is the obvious implementation.
The untangling is the part that can be wrong — a wrong twiddle or conjugate still returns
a plausible spectrum — so it is checked against a full fft of the same signal,
bin for bin, at every power of two to 256. Bin 0 and the Nyquist bin of a real signal are
themselves real.
try zn.convolve(f64, out, a, v, .full); // a.len + v.len - 1
try zn.convolve(f64, out, a, v, .same); // a.len
try zn.convolve(f64, out, a, v, .valid); // a.len - v.len + 1
try zn.correlate(f64, out, a, v, .full);convolve/// Discrete convolution of `a` with `v`, into `out`.
///
/// CONVOLUTION FLIPS THE KERNEL AND CORRELATION DOES NOT - which is the whole difference
/// between this and `correlate`, and the reason both exist. `conv2d` is named for the
/// framework convention and does not flip; this one is named for the mathematics and does.
pub fn convolve(
comptime T: type,
out: []T,
a: []const T,
v: []const T,
extent: ConvMode,
) Error!void {
comptime requireFloat(T);
return overlap(T, out, a, v, extent, true);
}
correlate/// Discrete cross-correlation of `a` with `v`, into `out`. The kernel is not flipped.
pub fn correlate(
comptime T: type,
out: []T,
a: []const T,
v: []const T,
extent: ConvMode,
) Error!void {
comptime requireFloat(T);
return overlap(T, out, a, v, extent, false);
}
Convolution flips the kernel and correlation does not — the whole
difference between the two, and the reason both exist. The test asserts
correlate(a, v) == convolve(a, reverse(v)) on an asymmetric kernel, and
separately that the two differ on it, so the identity has content. Note that
conv2d is named for the framework convention and does not flip; this one
is named for the mathematics and does.
Every mode is a window onto full and exactly that rather than
checking each mode's numbers independently — an off-by-one in same or
valid is the classic error, and a standalone test of either would not see it.
var file: std.ArrayList(u8) = .empty;
defer file.deinit(gpa);
const names = [_][]const u8{ "l1.w", "l1.b", "l2.w", "l2.b" };
try zn.saveTensors(f64, gpa, &file, &names, &.{ l1.weight, l1.bias, l2.weight, l2.bias });
try zn.loadTensors(f64, file.items, &names, &.{ r1.weight, r1.bias, r2.weight, r2.bias });saveTensorspub fn saveTensors(
comptime T: type,
gpa: Allocator,
out: *std.ArrayList(u8),
names: []const []const u8,
tensors: []const Tensor(T),
) Error!void {
comptime requireFloat(T);
if (names.len != tensors.len) {
return Error.ShapeMismatch;
}
for (tensors) |t| {
if (!t.isContiguous()) {
return Error.UnsupportedShape;
}
if (t.rank > max_rank) {
return Error.UnsupportedShape;
}
}
out.appendSlice(gpa, magic) catch return Error.OutOfMemory;
out.append(gpa, format_version) catch return Error.OutOfMemory;
try putInt(gpa, out, u32, @intCast(tensors.len));
for (names, tensors) |name, t| {
if (name.len > std.math.maxInt(u16)) {
return Error.UnsupportedShape;
}
try putInt(gpa, out, u16, @intCast(name.len));
out.appendSlice(gpa, name) catch return Error.OutOfMemory;
out.append(gpa, @backingInt(DType.of(T))) catch return Error.OutOfMemory;
out.append(gpa, t.rank) catch return Error.OutOfMemory;
var axis: usize = 0;
while (axis < t.rank) : (axis += 1) {
try putInt(gpa, out, u64, @intCast(t.shape[axis]));
}
const elements: []const T = t.data[t.base..][0..t.size()];
out.appendSlice(gpa, std.mem.sliceAsBytes(elements)) catch return Error.OutOfMemory;
}
}
loadTensors/// Read a file written by `saveTensors` into tensors the caller has already allocated, matched by
/// name. Every requested name must be present with the right dtype and shape.
///
/// INTO EXISTING TENSORS, NOT NEW ONES. Loading a checkpoint means filling the weights a model
/// already holds; allocating fresh tensors would leave the caller to copy them across and to free
/// two sets. It also means the shape check is free: the destination states the expected shape,
/// and a file that disagrees is `ShapeMismatch` rather than a silent reshape.
pub fn loadTensors(
comptime T: type,
bytes: []const u8,
names: []const []const u8,
into: []const Tensor(T),
) Error!void {
comptime requireFloat(T);
if (names.len != into.len) {
return Error.ShapeMismatch;
}
var cursor: usize = 0;
if (bytes.len < magic.len + 5 or !std.mem.eql(u8, bytes[0..magic.len], magic)) {
return Error.DomainError;
}
cursor = magic.len;
if (bytes[cursor] != format_version) {
return Error.DomainError;
}
cursor += 1;
const count: u32 = try takeInt(bytes, &cursor, u32);
var filled: usize = 0;
var entry: u32 = 0;
while (entry < count) : (entry += 1) {
const name_len: u16 = try takeInt(bytes, &cursor, u16);
if (cursor + name_len > bytes.len) {
return Error.DomainError;
}
const name: []const u8 = bytes[cursor..][0..name_len];
cursor += name_len;
if (cursor + 2 > bytes.len) {
return Error.DomainError;
}
const tag: u8 = bytes[cursor];
const rank: u8 = bytes[cursor + 1];
cursor += 2;
if (rank > max_rank) {
return Error.DomainError;
}
var shape: [max_rank]usize = @splat(1);
var elements: usize = 1;
var axis: usize = 0;
while (axis < rank) : (axis += 1) {
shape[axis] = @intCast(try takeInt(bytes, &cursor, u64));
elements *= shape[axis];
}
const byte_len: usize = elements * @sizeOf(T);
if (cursor + byte_len > bytes.len) {
return Error.DomainError;
}
const payload: []const u8 = bytes[cursor..][0..byte_len];
cursor += byte_len;
// Is this entry one the caller asked for?
for (names, into, 0..) |wanted, dest, i| {
if (!std.mem.eql(u8, wanted, name)) {
continue;
}
if (tag != @backingInt(DType.of(T))) {
return Error.DomainError;
}
if (dest.rank != rank or dest.size() != elements or !dest.isContiguous()) {
return Error.ShapeMismatch;
}
var d: usize = 0;
while (d < rank) : (d += 1) {
if (dest.shape[d] != shape[d]) {
return Error.ShapeMismatch;
}
}
@memcpy(std.mem.sliceAsBytes(dest.data[dest.base..][0..elements]), payload);
filled += 1;
_ = i;
break;
}
}
if (filled != names.len) {
return Error.DomainError;
}
}
magic 6 bytes "ZNUM\0\0"
version 1 byte
entry count 4 bytes little-endian
per entry:
name length 2 bytes
name n bytes
dtype tag 1 byte
rank 1 byte
shape 8 bytes per axis
element data sizeOf(T) per element, row-major, little-endianThe dtype tag encodes the type rather than counting it. The high nibble is the
kind (1 float, 2 signed, 3 unsigned) and the low nibble is log₂ of the byte width, so
f32 is 0x12 and f64 is 0x13. A reader computes
the element size straight from the tag — 1 << (tag & 0xf), which is
what byteWidth does — instead of consulting a table that has to be kept in step
with the enum, and an unknown tag still says how many bytes to skip. A dense counter would be
shorter to write and would make every future type a decision about where in the sequence it goes;
this way i32 is 0x22 by construction.
The test pins the header's actual bytes and asserts byteWidth against
@sizeOf for all three types, because a round trip agrees with itself whatever the
tags are — both sides of it share any mistake. Everything is little-endian at a stated width: @bitCasting a struct into a
file would be shorter and would break the first time the padding, the field order or the host's
endianness changed, none of which the file records. A reader in another language needs only the
table above.
loadTensors fills tensors the caller already holds, matched by name, rather than
allocating new ones. Loading a checkpoint means filling the weights a model already has;
allocating fresh ones would leave the caller to copy them across and free two sets. It also makes
the shape check free — the destination states the expected shape, and a file that disagrees
is ShapeMismatch rather than a silent reshape.
Every way the file can be wrong is refused rather than read as something plausible: a bad magic, an unknown version, a truncated file, a requested name the file does not hold, a shape that disagrees, and a dtype that does not match. The test asserts all seven, and asserts the round trip by the model's predictions rather than its weights — comparing weights would pass on a loader that read them in the wrong order into same-shaped tensors.
The source must be contiguous: the file stores elements in order, and a strided view's data is
not that order. materialise first.
try zn.saveNpy(f64, gpa, &bytes, a);
try zn.loadNpy(f64, out, bytes);
const header = try zn.readNpyHeader(bytes); // shape, width, where the data startssaveNpy/// Write one tensor in numpy's `.npy` format, version 1.0.
///
/// THE ONE FORMAT THAT LETS THIS LIBRARY TALK TO ANYTHING ELSE
///
/// `saveTensors` above is zimrnum's own container and is better at what it does - several named
/// tensors, a version byte, a self-describing dtype tag. But nothing else in the world reads it.
/// `.npy` is what numpy, scipy, PyTorch and every notebook already understand, and it costs
/// eighty lines.
///
/// The header is a Python dict literal in ASCII, padded with spaces so the DATA begins at a
/// 64-byte boundary. That padding is not decoration: it exists so a reader can memory-map the
/// array without copying, and writing it wrong produces a file every other tool rejects.
pub fn saveNpy(
comptime T: type,
gpa: Allocator,
out: *std.ArrayList(u8),
a: Tensor(T),
) Error!void {
comptime requireFloat(T);
var header: std.ArrayList(u8) = .empty;
defer header.deinit(gpa);
header.appendSlice(gpa, "{'descr': '") catch return Error.OutOfMemory;
header.appendSlice(gpa, comptime npyDescr(T)) catch return Error.OutOfMemory;
header.appendSlice(gpa, "', 'fortran_order': False, 'shape': (") catch return Error.OutOfMemory;
var digits: [24]u8 = undefined;
for (0..a.rank) |axis| {
const printed: []const u8 = bufPrint(&digits, "{d}", .{a.shape[axis]}) catch
return Error.OutOfMemory;
header.appendSlice(gpa, printed) catch return Error.OutOfMemory;
// A one-element tuple needs its trailing comma or Python reads it as a parenthesised
// number. numpy writes `(3,)` and expects to read it back.
if (a.rank == 1 or axis + 1 < a.rank) {
header.appendSlice(gpa, ",") catch return Error.OutOfMemory;
}
}
header.appendSlice(gpa, "), }") catch return Error.OutOfMemory;
// 10 bytes of preamble, then the header, then padding, then a newline - all so the data
// starts at a multiple of 64.
const preamble: usize = 10;
var padded: usize = header.items.len + 1;
while ((preamble + padded) % 64 != 0) {
padded += 1;
}
out.appendSlice(gpa, "\x93NUMPY") catch return Error.OutOfMemory;
out.append(gpa, 1) catch return Error.OutOfMemory;
out.append(gpa, 0) catch return Error.OutOfMemory;
const header_len: u16 = @intCast(padded);
out.append(gpa, @intCast(header_len & 0xFF)) catch return Error.OutOfMemory;
out.append(gpa, @intCast(header_len >> 8)) catch return Error.OutOfMemory;
out.appendSlice(gpa, header.items) catch return Error.OutOfMemory;
for (header.items.len..padded - 1) |_| {
out.append(gpa, ' ') catch return Error.OutOfMemory;
}
out.append(gpa, '\n') catch return Error.OutOfMemory;
var walk: Walk = .over(a.shape[0..a.rank]);
while (walk.next()) |at| {
const value: T = try a.at(at);
const bits: [@sizeOf(T)]u8 = @bitCast(value);
out.appendSlice(gpa, &bits) catch return Error.OutOfMemory;
}
}
loadNpy/// Read a `.npy` array into `out`, whose shape must match what the header says.
///
/// The width in the file need not match `T`: a `<f4` file read as f64 is widened per element,
/// which is exact. The other direction narrows and is allowed, because refusing it would mean a
/// caller who knows their data fits in f32 cannot say so.
pub fn loadNpy(comptime T: type, out: Tensor(T), bytes: []const u8) Error!void {
comptime requireFloat(T);
const header: NpyHeader = try readNpyHeader(bytes);
if (header.rank != out.rank) {
return Error.ShapeMismatch;
}
var count: usize = 1;
for (0..header.rank) |axis| {
if (header.shape[axis] != out.shape[axis]) {
return Error.ShapeMismatch;
}
count *= header.shape[axis];
}
if (bytes.len < header.data_at + count * header.width) {
return Error.DomainError;
}
const data: []const u8 = bytes[header.data_at..];
var index: usize = 0;
var walk: Walk = .over(out.shape[0..out.rank]);
while (walk.next()) |at| : (index += 1) {
const offset: usize = index * header.width;
const value: T = switch (header.width) {
4 => @floatCast(@as(f32, @bitCast(data[offset..][0..4].*))),
8 => @floatCast(@as(f64, @bitCast(data[offset..][0..8].*))),
else => unreachable, // readNpyHeader admits no other width
};
try out.setAt(at, value);
}
}
readNpyHeader/// Read a `.npy` header, refusing everything this library cannot represent.
///
/// WHAT IS REFUSED, AND WHY EACH ONE MATTERS
///
/// - **Big-endian** (`>f8`): the bytes would be read in the wrong order and produce plausible
/// nonsense rather than an error. Every value would be wrong and nothing would say so.
/// - **Fortran order**: the elements are all present and all in the wrong places. A reader that
/// ignored the flag would transpose the array silently.
/// - **Any dtype but `<f4` and `<f8`**: this library is float-only, and an integer array read as
/// floats is not a conversion, it is a reinterpretation.
///
/// Each of these is a case where CARRYING ON produces a result rather than a failure, which is
/// exactly when an explicit refusal is worth its lines.
pub fn readNpyHeader(bytes: []const u8) Error!NpyHeader {
if (bytes.len < 10 or !std.mem.eql(u8, bytes[0..6], "\x93NUMPY")) {
return Error.DomainError;
}
if (bytes[6] != 1) {
return Error.DomainError; // version 2 and 3 use a wider length field
}
const header_len: usize = @as(usize, bytes[8]) | (@as(usize, bytes[9]) << 8);
const header_end: usize = 10 + header_len;
if (header_end > bytes.len) {
return Error.DomainError;
}
const header: []const u8 = bytes[10..header_end];
const width: u8 = if (std.mem.indexOf(u8, header, "'<f8'") != null)
8
else if (std.mem.indexOf(u8, header, "'<f4'") != null)
4
else
return Error.UnsupportedShape;
if (std.mem.indexOf(u8, header, "'fortran_order': False") == null) {
return Error.UnsupportedShape;
}
const open_paren: usize = std.mem.indexOfScalar(u8, header, '(') orelse
return Error.DomainError;
const close_paren: usize = std.mem.indexOfScalarPos(u8, header, open_paren, ')') orelse
return Error.DomainError;
var out: NpyHeader = .{
.shape = @splat(1),
.rank = 0,
.data_at = header_end,
.width = width,
};
var cursor: usize = open_paren + 1;
while (cursor < close_paren) {
while (cursor < close_paren and !std.ascii.isDigit(header[cursor])) {
cursor += 1;
}
if (cursor >= close_paren) {
break;
}
var extent: usize = 0;
while (cursor < close_paren and std.ascii.isDigit(header[cursor])) {
extent = extent * 10 + (header[cursor] - '0');
cursor += 1;
}
if (out.rank >= max_rank) {
return Error.UnsupportedShape;
}
out.shape[out.rank] = extent;
out.rank += 1;
}
return out;
}
saveTensors is zimrnum's own container and is better at what it does —
several named tensors, a version byte, a self-describing dtype tag. But nothing else in
the world reads it. .npy is what numpy, scipy, PyTorch and every notebook
already understand, and it costs eighty lines.
The header is a Python dict literal in ASCII, padded with spaces so the data begins at a 64-byte boundary. That padding is not decoration: it exists so a reader can memory-map the array without copying, and getting it wrong produces a file that this library still reads and every other tool rejects. The test asserts the boundary directly, because a round trip through one implementation would never catch it.
Verified against real numpy, not just against itself. Files numpy wrote were
read here — including a <f4 widened to f64, and a (1,) shape
whose header carries the trailing comma Python needs — and bytes written here were handed
back to numpy, which read them as a C-contiguous (3, 4) float64 array with the right values.
Three things are refused, and each is a case where carrying on produces a result rather than a failure:
| Refused | What would happen otherwise |
|---|---|
big-endian >f8 | every value read in the wrong byte order — plausible nonsense, no error |
fortran_order: True | every element present, every one in the wrong place: a silent transpose |
any dtype but <f4/<f8 | an integer array read as floats is not a conversion, it is a reinterpretation |
try zn.generalizedAdvantage(f64, rewards, values, next_values, terminal, episode_end, 0.99, 0.95, adv, ret);
try zn.discountedReturns(f64, rewards, terminal, 0.99, out);
try zn.normalizeAdvantages(f64, adv, 1.0e-8);
const objective = try zn.ppoClipObjective(f64, logp_new, logp_old, adv, 0.2);
try zn.entropyRows(f64, out, probs);
var ring: zn.ReplayBuffer(f32) = try .init(gpa, 100_000, width);
try ring.push(transition);
try ring.sample(rng, index, batch);generalizedAdvantage/// Generalised advantage estimation over one rollout, walked backwards.
///
/// TWO KINDS OF ENDING, AND THEY ARE NOT THE SAME
///
/// `terminal[i]` is 1 where the episode truly ended - the pole fell, the agent died. No value
/// lies beyond it, so nothing is bootstrapped. `episode_end[i]` is 1 where the ROLLOUT was cut -
/// a time limit, a batch boundary - while the episode itself continues: the next state's value is
/// still bootstrapped, but the running advantage sum is reset so it does not leak across the cut.
/// Conflating the two is the most common GAE bug and it produces an agent that learns, slowly,
/// from returns that are wrong at every episode boundary.
///
/// This follows znum's form exactly: `bootstrap = gamma * next_value * (1 - terminal)`,
/// `delta = r + bootstrap - value`, `running = delta + gammalambda * (1 - episode_end) * running`.
///
/// `out_returns = advantage + value` is the target for the value function, which is why both
/// come back from one pass.
pub fn generalizedAdvantage(
comptime T: type,
rewards: []const T,
values: []const T,
next_values: []const T,
terminal: []const T,
episode_end: []const T,
gamma: T,
lambda: T,
out_advantages: []T,
out_returns: []T,
) Error!void {
comptime requireFloat(T);
const n: usize = rewards.len;
if (values.len != n or next_values.len != n or terminal.len != n or
episode_end.len != n or out_advantages.len != n or out_returns.len != n)
{
return Error.ShapeMismatch;
}
var running: T = 0;
var i: usize = n;
while (i > 0) {
i -= 1;
const bootstrap: T = gamma * next_values[i] * (1 - terminal[i]);
const delta: T = rewards[i] + bootstrap - values[i];
running = delta + gamma * lambda * (1 - episode_end[i]) * running;
out_advantages[i] = running;
out_returns[i] = running + values[i];
}
}
discountedReturns/// Discounted returns, walked backwards, reset at each terminal.
///
/// GAE with `lambda = 1` and zero values reduces to this, and the test asserts it: an
/// implementation of one that disagreed with the other would be wrong in a way a single test
/// of either cannot see.
pub fn discountedReturns(
comptime T: type,
rewards: []const T,
terminal: []const T,
gamma: T,
out: []T,
) Error!void {
comptime requireFloat(T);
if (terminal.len != rewards.len or out.len != rewards.len) {
return Error.ShapeMismatch;
}
var running: T = 0;
var i: usize = rewards.len;
while (i > 0) {
i -= 1;
running = rewards[i] + gamma * (1 - terminal[i]) * running;
out[i] = running;
}
}
normalizeAdvantages/// Standardise advantages in place: zero mean, unit deviation, with a small epsilon.
///
/// Sample deviation, and an epsilon rather than a `DomainError` for a constant batch: an
/// all-equal advantage batch is a legitimate (if uninformative) rollout, not a caller's mistake,
/// and PPO must keep stepping through it.
pub fn normalizeAdvantages(comptime T: type, advantages: []T, epsilon: T) Error!void {
comptime requireFloat(T);
if (advantages.len < 2) {
return Error.DomainError;
}
const t: Tensor(T) = try Tensor(T).fromSlice(advantages, &.{advantages.len});
const mean: T = try meanAll(T, t);
const sd: T = try stdDev(T, t, .sample);
for (advantages) |*a| {
a.* = (a.* - mean) / (sd + epsilon);
}
}
ppoClipObjective/// PPO's clipped surrogate objective, as a scalar to MAXIMISE. Returns the mean over the batch.
///
/// THE MINIMUM IS THE WHOLE IDEA
///
/// `ratio = exp(logp_new - logp_old)`. The unclipped term `ratio * A` rewards moving the policy
/// toward actions with positive advantage without limit; the clipped term holds the ratio to
/// `[1 - epsilon, 1 + epsilon]`. Taking the MINIMUM of the two is pessimistic in exactly the right way: for a
/// positive advantage the objective stops improving once the ratio passes `1 + epsilon`, and for a
/// negative one once it drops below `1 - epsilon`. There is no incentive to move the policy further
/// than the clip, which is what keeps a PPO update from destroying a policy in one step.
///
/// Follows znum's `clippedSurrogate` term for term. Returned as an objective rather than a loss;
/// negate it for a gradient step.
pub fn ppoClipObjective(
comptime T: type,
logp_new: []const T,
logp_old: []const T,
advantages: []const T,
clip: T,
) Error!T {
comptime requireFloat(T);
const n: usize = logp_new.len;
if (logp_old.len != n or advantages.len != n) {
return Error.ShapeMismatch;
}
if (n == 0 or clip <= 0) {
return Error.DomainError;
}
var total: T = 0;
for (logp_new, logp_old, advantages) |ln, lo, a| {
const ratio: T = @exp(ln - lo);
const unclipped: T = ratio * a;
const clipped: T = scalarClamp(ratio, 1 - clip, 1 + clip) * a;
total += @min(unclipped, clipped);
}
return total / @as(T, @floatFromInt(n));
}
entropyRows/// The entropy of each row of a rank-2 tensor of probabilities, into `out`.
///
/// A zero probability contributes zero, by the limit of `p*log p`; a uniform row gives
/// `log(width)`, which is the maximum and is asserted by the test.
pub fn entropyRows(comptime T: type, out: []T, probs: Tensor(T)) Error!void {
comptime requireFloat(T);
if (probs.rank != 2) {
return Error.UnsupportedShape;
}
if (out.len != probs.shape[0]) {
return Error.ShapeMismatch;
}
for (out, 0..) |*slot, row| {
var h: T = 0;
var col: usize = 0;
while (col < probs.shape[1]) : (col += 1) {
const p: T = try probs.at(&.{ row, col });
if (p > 0) {
h -= p * @log(p);
}
}
slot.* = h;
}
}
ReplayBuffer/// A ring buffer of fixed-width transitions with uniform random sampling.
///
/// A RING, SO OLD EXPERIENCE FALLS OUT
///
/// Off-policy methods learn from a replay of past transitions, and the buffer must forget: a
/// transition gathered by a policy from ten thousand steps ago says less about the current one
/// than a recent transition does. Overwriting the oldest slot is the simplest schedule that
/// forgets, and it needs no bookkeeping beyond a write index.
///
/// A transition is a flat `[]T` of `width` values - the caller decides the layout
/// (`state ++ action ++ reward ++ next_state ++ done` is usual). The buffer does not interpret
/// it, which is what lets one type serve discrete and continuous problems alike.
pub fn ReplayBuffer(comptime T: type) type {
comptime requireFloat(T);
return struct {
const Self = @This();
data: []T,
width: usize,
capacity: usize,
/// Next slot to write. Wraps.
head: usize,
/// How many slots hold a transition; saturates at `capacity`.
len: usize,
pub fn init(gpa: Allocator, capacity: usize, width: usize) Error!Self {
if (capacity == 0 or width == 0) {
return Error.DomainError;
}
return .{
.data = try gpa.alloc(T, capacity * width),
.width = width,
.capacity = capacity,
.head = 0,
.len = 0,
};
}
/// Store one transition, overwriting the oldest once full.
pub fn push(self: *Self, transition: []const T) Error!void {
if (transition.len != self.width) {
return Error.ShapeMismatch;
}
@memcpy(self.data[self.head * self.width ..][0..self.width], transition);
self.head = (self.head + 1) % self.capacity;
self.len = @min(self.len + 1, self.capacity);
}
/// Copy `out.len / width` uniformly random transitions into `out`, with replacement.
/// `index` addresses the draw, so the same `(rng, index)` samples the same batch.
pub fn sample(self: *const Self, rng: Rng, index: u32, out: []T) Error!void {
if (self.len == 0 or out.len % self.width != 0) {
return Error.DomainError;
}
const batch: usize = out.len / self.width;
var i: usize = 0;
while (i < batch) : (i += 1) {
const slot: usize = rng.intBelow(index *% 0x9e37 +% @as(u32, @intCast(i)), @intCast(self.len));
@memcpy(out[i * self.width ..][0..self.width], self.data[slot * self.width ..][0..self.width]);
}
}
};
}
generalizedAdvantage takes two flags per step. terminal is 1 where the episode truly
ended: no value lies beyond it, so nothing is bootstrapped. episode_end is 1 where
the rollout was cut — a time limit, a batch boundary — while the episode
continues: the next state's value is still bootstrapped, but the running sum resets so it does
not leak across the cut.
bootstrap = γ · next_value[i] · (1 − terminal[i])
δ = reward[i] + bootstrap − value[i]
running = δ + γλ · (1 − episode_end[i]) · running // walked backwards
advantage[i] = running
return[i] = running + value[i]Conflating the two flags is the most common GAE error, and it produces an agent that learns,
slowly, from returns that are wrong at every episode boundary. The test sets a true terminal at
one step and a truncation at another and checks that only the truncation bootstraps.
discountedReturns is what generalizedAdvantage reduces to with λ = 1 and zero
values, and the two are asserted equal.
ratio = exp(logp_new − logp_old)
unclipped = ratio · A
clipped = clamp(ratio, 1 − ε, 1 + ε) · A
objective = mean(min(unclipped, clipped))For a positive advantage the objective stops improving once the ratio passes
1 + ε; for a negative one, once it drops below 1 − ε.
There is no incentive to move the policy further than the clip, which is what keeps one update
from destroying a policy. Moving the wrong way is never clipped away: a ratio of 5 with a
negative advantage is penalised in full. It is returned as an objective to maximise; negate it for
a gradient step.
ReplayBuffer stores fixed-width transitions and overwrites the oldest once full.
A transition gathered ten thousand steps ago says less about the current policy than a recent
one, and overwriting the oldest slot is the simplest schedule that forgets. A transition is a
flat slice whose layout the caller decides — the buffer does not interpret it, which lets one
type serve discrete and continuous problems. sample is addressable: the same
(rng, index) yields the same batch.
var s = zn.cartpoleReset(f64, rng, episode);
const step = zn.cartpoleStep(f64, s, .right);
// step.state, step.reward, step.failedcartpoleReset/// A fresh cartpole, every value drawn uniformly from a small band around upright.
///
/// The band is the standard one. Starting exactly upright would make the first few steps
/// deterministic and let a policy memorise them.
pub fn cartpoleReset(comptime T: type, rng: Rng, episode: u32) CartpoleState(T) {
comptime requireFloat(T);
const band: T = 0.05;
const stream: Rng = rng.split(episode);
return .{
.cart = (stream.unitFloat(T, 0) - 0.5) * 2 * band,
.cart_rate = (stream.unitFloat(T, 1) - 0.5) * 2 * band,
// The band is the standard one and is quoted in radians, like the pole's angle.
.pole_rad = (stream.unitFloat(T, 2) - 0.5) * 2 * band,
.pole_rate_rad = (stream.unitFloat(T, 3) - 0.5) * 2 * band,
};
}
cartpoleStep/// The cart-and-pole control benchmark: push a cart left or right to keep a pole upright.
///
/// A PURE FUNCTION - THE STATE GOES IN AND THE NEXT STATE COMES OUT
///
/// znum's step takes `state: *Tensor(f32)` and updates it IN PLACE, with a doc comment warning
/// that the tensor "holds the observation for step t only until this returns, so a rollout copies
/// it out BEFORE calling". **That is a footgun stated in prose**: the caller has to remember, and
/// forgetting produces a rollout whose observations are all one step late, which trains and
/// converges to something slightly wrong.
///
/// Returning the next state makes it unforgettable. It also makes the step replayable - the same
/// state and action give the same answer, always - which is what a test needs and what a
/// debugger needs.
///
/// RESET IS A SEPARATE FUNCTION, AND THAT REMOVES TWO FLAGS
///
/// znum folds the reset into the step, so the step needs a `force_reset` parameter, and then a
/// `no_transition` special case because a forced reset produces no reward and no termination.
/// Both disappear when starting an episode is its own call: `cartpoleReset` gives a state, and
/// `cartpoleStep` only ever steps.
///
/// RADIANS HERE, AND THE SUFFIX SAYS SO - BECAUSE THIS ANGLE IS DIFFERENTIATED
///
/// The rest of this codebase moved to turns, and for geometry that was right: a quarter turn is
/// exactly 0.25, and `sinTurns` is exact where `sin(tau/4)` is not.
///
/// **A differentiated angle is the other case.** `d/dtheta sin(theta)` is `cos(theta)` only in
/// radians; in turns it is `tau * cos`, and every derivative in a dynamics system picks up that
/// factor. The cartpole equations are differential equations, so radians are the unit in which
/// they have no constants in them at all.
///
/// The first version of this held the pole in turns and converted the angular acceleration from
/// radians on the way in. That conversion was a place to be wrong, and I was: **the pole fell tau
/// times too fast** and every sign still looked right. Choosing the unit the equations are
/// written in removes the conversion rather than fixing it.
///
/// So: **turns for angles that are drawn, radians for angles that are differentiated**, and the
/// name carries the unit either way. `pole_rad` cannot be mistaken for `pole_turns`.
///
/// The limit is still not a rounded constant. znum writes `0.2095`, which is 12.0035 degrees;
/// `radFromDeg(12)` is twelve degrees exactly as far as the type can express it.
///
/// The dynamics are the standard ones - Barto, Sutton and Anderson 1983 - and the integration is
/// SEMI-IMPLICIT Euler: the position update uses the NEW velocity. Gym's default is explicit
/// Euler, which drifts; znum uses semi-implicit too and is right to.
pub fn cartpoleStep(
comptime T: type,
state: CartpoleState(T),
push: Push,
) CartpoleStep(T) {
comptime requireFloat(T);
const gravity: T = 9.8;
const cart_mass: T = 1.0;
const pole_mass: T = 0.1;
const total_mass: T = cart_mass + pole_mass;
const half_pole: T = 0.5;
const pole_moment: T = pole_mass * half_pole;
const push_force: T = 10.0;
const seconds: T = 0.02;
const force: T = switch (push) {
.left => -push_force,
.right => push_force,
};
// `zm.` because zimrnum has its own `sinRad` and it is a tensor op - the disambiguation the
// two `lint:off` lines at the top of this file exist for.
const lean: T = zm.sinRad(state.pole_rad);
const upright: T = zm.cosRad(state.pole_rad);
// NO CONVERSION ANYWHERE, WHICH IS THE POINT OF CHOOSING RADIANS
//
// These are the textbook equations transcribed, and every quantity is already in the unit the
// derivation uses. An earlier version held the pole in turns and converted here; that line
// was a place to be wrong and it was.
const shared: T = (force + pole_moment * state.pole_rate_rad * state.pole_rate_rad * lean) /
total_mass;
const pole_accel: T = (gravity * lean - upright * shared) /
(half_pole * (4.0 / 3.0 - pole_mass * upright * upright / total_mass));
const cart_accel: T = shared - pole_moment * pole_accel * upright / total_mass;
// Semi-implicit: the position uses the UPDATED velocity, which conserves energy far better
// than the explicit form over a long episode.
const cart_rate: T = state.cart_rate + seconds * cart_accel;
const pole_rate_rad: T = state.pole_rate_rad + seconds * pole_accel;
const next: CartpoleState(T) = .{
.cart = state.cart + seconds * cart_rate,
.cart_rate = cart_rate,
.pole_rad = state.pole_rad + seconds * pole_rate_rad,
.pole_rate_rad = pole_rate_rad,
};
// Twelve degrees, converted rather than rounded - and 2.4 metres of track either side.
const fall_rad: T = zm.radFromDeg(@as(T, 12.0));
const track: T = 2.4;
const failed: bool = @abs(next.cart) > track or @abs(next.pole_rad) > fall_rad;
return .{
.state = next,
// A reward of one for every step survived, which is the standard scoring and makes the
// return equal to the episode length - so "did it learn" reads directly off the number.
.reward = 1.0,
.failed = failed,
};
}
The step is pure. znum's takes state: *Tensor(f32) and updates it
in place, with a doc comment warning that the tensor "holds the observation for step t only until
this returns, so a rollout copies it out BEFORE calling". That is a footgun stated in prose —
forgetting produces a rollout whose observations are all one step late, which trains and converges
to something slightly wrong.
Reset is a separate call, and that removes two things. znum folds the reset into
the step, so the step needs a force_reset parameter and then a
no_transition special case because a forced reset produces no reward and no
termination. Both disappear when starting an episode is its own function.
The four numbers are named. znum indexes state_data[base + 2] for
the pole angle; nothing at the call site says which of four slots holds what, and swapping two
produces a plausible simulation that is wrong. Same for rd[0] and rd[1],
which are here reward and failed.
The rest of this library moved to turns, and for geometry that was right: a quarter turn is
exactly 0.25, and sinTurns is exact where sin(tau/4) is not.
A differentiated angle is the other case.
d/dθ sin(θ) is cos(θ) only in radians; in turns
it is tau · cos, and every derivative in a dynamics system picks up that factor.
The cartpole equations are differential equations, so radians are the unit in which they contain no
constants at all.
A version held the pole in turns and converted the angular acceleration on the way in. That conversion was a place to be wrong, and it was: the pole fell tau times too fast. Every sign was right, the episode ended, the reward accumulated — it looked exactly like a cartpole, and a test checking only directions passed. A unit error inside a plausible simulation has no symptom.
Choosing the unit the equations are written in removed the conversion rather than
fixing it. So the rule is decided per piece of angle state, and the name always carries the unit:
pole_rad and pole_rate_rad cannot be mistaken for
pole_turns.
The test still asserts a magnitude rather than a direction — 9.0977 rad/s² over 0.02 s is 0.295821 rad/s — but that is now the number the equations give directly, with nothing converted on the way.
The limit is not a rounded constant either. znum writes 0.2095, which is 12.0035
degrees; radFromDeg(12) is twelve degrees as exactly as the type can hold.
The integration is semi-implicit Euler: the position update uses the new velocity. Gym's default is explicit Euler, which drifts over a long episode; znum uses semi-implicit too and is right to.
Every piece has its own test. cartpoleStep has the right physics,
categorical has the right distribution, Dense has a checked gradient, the
tape matches finite differences. None of that says an agent can learn.
A sign error in the policy update, an observation one step stale, a baseline drifting the wrong way — each of those leaves every component test green and the agent at chance. So the suite trains one: cartpole, a two-layer policy, REINFORCE with a running baseline, forty rounds of eight episodes.
seed first round last round
1 16.0 40.8
2 21.0 44.3
3 13.4 20.4
4 15.3 35.4
5 23.4 47.9Five seeds measured before choosing the bar, which is a 40% rise — under the worst seed's margin so it will not flake, and far above the noise of a policy that is not learning. An untrained cartpole survives about twenty steps, so that is the floor to beat.
It is deliberately the cheapest thing that can fail honestly: no target network, no value head, no clipping. The sign of the advantage is the entire algorithm — backwards, it learns to fail, and every other test stays green.
ppoClipObjective returns a scalar, not a Var. It computes
PPO's formula but is not on the tape, so it cannot drive a gradient — the library had PPO's
arithmetic and not PPO's training step. That gap was invisible until something tried to train with
it, which is exactly the kind of thing an end-to-end test is for.
graph.ppoClipLoss closes it, and the interesting part is where the
clip lives.
const loss = try graph.ppoClipLoss(logp_new, logp_old, advantages, 0.2);
try graph.backward(loss);In the forward the clip is a min, which looks like a detail. In the backward
it is the algorithm: where the clipped branch is smaller, the objective no longer depends
on the ratio, so the gradient is exactly zero and the update stops pushing a ratio that has
already moved too far. That is what keeps a policy step from destroying the policy, and a version
computing only the right loss would still diverge.
The test checks the gradient against central differences at every element — a piecewise objective produces plausible numbers on the wrong branch — and then asserts the zeros directly: one index clipped from above with a positive advantage, one clipped from below with a negative one, and one unclipped index that must not be zero, or the check would pass against a backward that returns zero everywhere.
logp_old and advantages are plain slices rather than
Vars, because both are constants of the step: the old policy is fixed by definition and
the advantages were estimated before it. They ride on the node in a constants field
rather than overloading labels, which is []const usize and says "class
indices" — a slot whose name is a lie costs more than a field.
try zn.gae(steps, 0.99, 0.95, &advantages, &returns);gae/// Generalized Advantage Estimation over a rollout, plus the returns to fit a critic against.
///
/// WHAT IT IS FOR
///
/// A policy gradient needs to know whether an action did better than expected. The crude answer
/// is the whole discounted return, which is unbiased and extremely noisy; the other extreme is
/// the one-step TD error, which is low-variance and biased by whatever the critic gets wrong.
///
/// GAE interpolates with `lambda`: **0 gives the one-step error, 1 gives the full return**, and
/// values between trade bias for variance. 0.95 is the usual choice and is not magic.
///
/// `gamma` discounts the future as everywhere else in RL.
///
/// THE RECURRENCE RUNS BACKWARD, AND IT HAS TO
///
/// Each advantage depends on the one after it, so a forward loop would need the answer before it
/// had it. Walking back from the end costs one pass and no extra storage.
pub fn gae(
steps: []const RolloutStep,
gamma: f64,
lambda: f64,
out_advantages: []f64,
out_returns: []f64,
) Error!void {
if (out_advantages.len != steps.len or out_returns.len != steps.len) {
return Error.ShapeMismatch;
}
if (gamma < 0 or gamma > 1 or lambda < 0 or lambda > 1) {
return Error.DomainError;
}
var running: f64 = 0;
var i: usize = steps.len;
while (i > 0) {
i -= 1;
const step: RolloutStep = steps[i];
// `terminal` kills the BOOTSTRAP: no next state exists, so its value is zero.
const bootstrap: f64 = if (step.terminal) 0 else gamma * step.next_value;
const delta: f64 = step.reward + bootstrap - step.value;
// `episode_end` kills the TRACE: the next step's advantage belongs to another trajectory.
running = if (step.episode_end) delta else delta + gamma * lambda * running;
out_advantages[i] = running;
// The return the critic should have predicted, which is the advantage plus what it did
// predict. Fitting the critic against this rather than the raw discounted return is what
// makes the two halves consistent.
out_returns[i] = running + step.value;
}
}
A policy gradient needs to know whether an action did better than expected. The crude answer is the whole discounted return: unbiased and extremely noisy. The other extreme is the one-step TD error: low variance, biased by whatever the critic gets wrong.
GAE interpolates with lambda:
lambda 0 the one-step TD error
lambda 0.5 between
lambda 1 the full discounted return0.95 is the usual choice and is not magic. gamma discounts the future as
everywhere else.
These are two different things and treating them as one flag is the classic mistake.
terminal means the episode genuinely ended — the pole fell. There is no
future to bootstrap from, so the next state's value is zero by definition, not whatever
the critic happens to predict about a state that does not exist.
episode_end means the rollout stopped here, terminal or not. What stops is the
trace: an advantage must not carry across a seam where the following reward belongs to a
different trajectory.
reward 1, value 1, next_value 99
terminal = true -> advantage 0 the 99 is ignored
terminal = false -> advantage 99 the 99 is usedThe case one flag cannot express is a truncation: the step limit was reached
while the episode was still running. The next state exists and its predicted value is real, so the
bootstrap stands — but the trace must still stop. Using terminal for that throws
away a legitimate estimate and biases the critic low; using episode_end alone for a
real termination bootstraps off a state the agent never reached and biases it high.
Each advantage depends on the one after it, so a forward loop would need the answer before it had it. Walking back from the end costs one pass and no extra storage.
returns comes out as the advantage plus the critic's own prediction. Fitting the
critic against that rather than the raw discounted return is what keeps the two halves
consistent.
const policy = zn.Categorical(f64){ .logits = scores };
const action = try policy.sample(rng, row, stream);
const lp = try policy.logProb(gpa, row, action);
const bonus = try policy.entropy(gpa, row);Categorical/// Write a network's parameters, named by position.
///
/// THE GAP THIS CLOSES: A TRAINED MODEL YOU CANNOT SAVE IS A DEMO
///
/// `saveTensors` and `loadTensors` have been here all along, and they take names and tensors.
/// What was missing is the bridge from a network to those two lists - so a library that could
/// train a cartpole agent could not keep it, and every run started from noise.
///
/// NAMED BY POSITION, AND THAT IS A REAL CONSTRAINT
///
/// The names are `p0`, `p1`, ... in `attach` order. A `Var` carries no name, so there is nothing
/// better available - and the consequence has to be stated rather than discovered: **the loading
/// network must have the same architecture in the same order.** `loadTensors` checks the SHAPES,
/// so a mismatched layer is caught; a same-shaped layer in a different role is not.
///
/// znum's `StateDict` carries owned names, which is better for a network assembled dynamically.
/// This is the honest version for a `Chain`, whose layers are a comptime tuple and therefore
/// always in the same order.
/// A distribution over discrete actions, parameterised by logits.
///
/// WHY THIS IS A TYPE AND NOT THREE FUNCTIONS
///
/// A policy gradient needs `log p(action)` for the action it actually took, and an exploration
/// bonus needs the entropy. Both have to be computed from **the same normalisation `sample`
/// used** - otherwise the gradient points somewhere the sampler never went.
///
/// Bundling them means the three cannot disagree. They are separate functions in znum and the
/// caller is responsible for pairing them correctly.
///
/// LOGITS, NOT PROBABILITIES
///
/// The input is unnormalised scores. A network's last layer produces exactly that, and asking a
/// caller for probabilities means asking them to normalise - which is the step that overflows
/// when a logit is large, and which `logSoftmaxRows` exists to do stably.
pub fn Categorical(comptime T: type) type {
comptime requireFloat(T);
return struct {
const Self = @This();
/// Unnormalised scores, one row per batch element.
logits: Tensor(T),
/// Draw one action for row `row`.
///
/// Delegates to `categorical`, which subtracts the row maximum before exponentiating and
/// walks a cumulative sum. `index` selects an independent stream so two draws in one
/// step do not correlate.
pub fn sample(self: Self, rng: Rng, row: usize, index: u32) Error!usize {
if (self.logits.rank != 2 or row >= self.logits.shape[0]) {
return Error.OutOfRange;
}
const slice: Tensor(T) = try (try self.logits.slice(0, row, 1)).reshape(
&.{self.logits.shape[1]},
);
return categorical(T, slice, rng, index);
}
/// `log p(action)` for one row.
///
/// This is `logSoftmax(logits)[action]`, computed by the stable route rather than as
/// `log(softmax(x))`. The naive form loses the value entirely once a probability falls
/// below the smallest representable float - and for a confident policy that is the
/// common case, not the edge case.
pub fn logProb(self: Self, gpa: Allocator, row: usize, action: usize) Error!T {
const classes: usize = self.logits.shape[1];
if (self.logits.rank != 2 or row >= self.logits.shape[0] or action >= classes) {
return Error.OutOfRange;
}
const normalised: Tensor(T) = try Tensor(T).alloc(gpa, &.{ 1, classes });
defer gpa.free(normalised.data);
const one_row: Tensor(T) = try self.logits.slice(0, row, 1);
try logSoftmaxRows(T, normalised, one_row);
return normalised.at2(0, action);
}
/// The Shannon entropy of one row, in nats.
///
/// High entropy means the policy is still exploring; zero means it always picks the same
/// action. Added to a loss with a small coefficient, it is what stops a policy collapsing
/// onto one action before it has tried the others.
pub fn entropy(self: Self, gpa: Allocator, row: usize) Error!T {
const classes: usize = self.logits.shape[1];
if (self.logits.rank != 2 or row >= self.logits.shape[0]) {
return Error.OutOfRange;
}
const normalised: Tensor(T) = try Tensor(T).alloc(gpa, &.{ 1, classes });
defer gpa.free(normalised.data);
try logSoftmaxRows(T, normalised, try self.logits.slice(0, row, 1));
var total: CompensatedSum(T) = .{};
for (0..classes) |k| {
const log_p: T = normalised.at2(0, k);
// `p * log p` from the LOG, so a vanishing probability contributes a vanishing
// term rather than `0 * -inf`, which is NaN.
total.add(@exp(log_p) * log_p);
}
return -total.value();
}
};
}
A policy gradient needs log p(action) for the action actually taken, and an
exploration bonus needs the entropy. Both must be computed from the same normalisation
sample used — otherwise the gradient points somewhere the sampler never
went.
Bundling them means the three cannot disagree.
The input is unnormalised scores, which is what a network's last layer produces. Asking for probabilities means asking the caller to normalise, and that is the step that loses the value:
logits 0 and 800, asking for log p of the first
log(softmax(x)) -> -inf softmax gave exactly 0
logSoftmax(x) -> -800 correctFor a confident policy that is the common case, not the edge case — and a policy
gradient that reads -inf has nothing to learn from.
It is highest when the policy is uniform: over three actions, exactly log 3. Added
to a loss with a small coefficient, it rewards a policy for not yet having made up its mind, which
is what stops a collapse onto one action before the others have been tried.
The term p log p is computed from the log, not the probability. From the
probability a vanishing action would give 0 × -inf, which is NaN; from the log it
vanishes.
A network outputs an unconstrained real, and a standard deviation must be positive. Taking
exp guarantees that for free. Parameterising the deviation directly needs a clamp, and
a clamp that ever binds gives a zero gradient on the parameter it clamped —
so the network cannot learn its way out.
It also spaces the parameter sensibly: an equal step in log_std is an equal
ratio change in the spread, which is how scale behaves.
"Diagonal" records that the dimensions are independent, so the joint log density is the sum of the per-dimension ones. A full covariance would need the off-diagonal terms and a matrix inverse.
A Gaussian's entropy does not depend on its mean. Shifting a distribution moves
it without changing its shape, and entropy measures shape. Only the spread matters:
0.5 log(2πe) per dimension, plus log_std.
var pre: [dim]f64 = undefined;
var action: [dim]f64 = undefined;
try policy.sample(rng, row, &pre, &action); // keep BOTH
const lp = try policy.logProb(row, &pre); // scored from `pre`, not `action`A continuous action usually has bounds — a motor torque, a steering angle. An unbounded Gaussian will sometimes propose a value outside them, and clipping afterwards is not free: the clipped action has a different density from the one the policy scored, so the gradient no longer matches what happened. Squashing puts the bound inside the distribution.
Transforming a random variable changes its density by the derivative of the transform:
log p(a) = log p(u) - log(1 - tanh(u)^2) with a = tanh(u)Written that way it is useless in floating point:
u = 20 -> tanh(u) == 1.0 exactly
1 - 1 = 0
log(0) = -infAnd a trained agent lives exactly there, saturated against its action bound. The stable identity instead:
-log(1 - tanh(u)^2) = -2 log2 + 2u + 2 softplus(-2u)Every term stays finite for any u, because softplus is built not to
form the 1. This is the same trick logSumExp and log(softmax) need, in a
third place.
Not the action. Recovering u from a would need atanh,
which is unbounded as the action approaches its limits and loses all precision in f32 exactly where
a saturated policy sits.
Storing u alongside the action costs one float per dimension and removes the
problem, so sample returns both and the caller keeps both.
log_std is clamped to [-20, 2] — SAC's usual values — and
the clamp is applied inside both sample and logProb from the same
two constants.
That is the whole reason they are public constants on the type. A policy that samples
with one spread and scores with another is off-policy against itself, and the importance
ratio it computes means nothing. The clamp matters in practice because an unclamped head emits a
log_std of −60 early in training, whose exp underflows to a zero
spread and whose log-density is then infinite.
const stats = try zn.ppoObjective(
advantages, // modified in place if normalising
log_prob_old,
log_prob_new,
entropies,
.{}, // PpoConfig defaults
);ppoObjective/// The clipped surrogate objective over a batch, with the diagnostics that go with it.
///
/// Returns the loss to MINIMISE: the negated objective, so it drops into a gradient step like
/// any other loss. `advantages` is modified in place when the config asks for normalisation.
///
/// THE KL ESTIMATOR IS `exp(r) - 1 - r`, NOT `old - new`
///
/// The cheap estimator every reference implementation started with is `mean(log_old - log_new)`,
/// which is unbiased and **can come out negative** - and a negative divergence is not a small
/// number, it is a wrong one. A caller thresholding on it either never triggers or triggers on
/// noise.
///
/// The k3 estimator `exp(r) - 1 - r` with `r = log_new - log_old` is **non-negative for every
/// r**, because `exp(r) >= 1 + r` everywhere. znum uses the same one.
pub fn ppoObjective(
advantages: []f64,
log_prob_old: []const f64,
log_prob_new: []const f64,
entropies: []const f64,
config: PpoConfig,
) Error!PpoStats {
const n: usize = advantages.len;
if (log_prob_old.len != n or log_prob_new.len != n or entropies.len != n) {
return Error.ShapeMismatch;
}
if (n == 0) {
return Error.DomainError;
}
if (config.clip <= 0) {
return Error.DomainError;
}
if (config.normalize_advantage) {
var total: CompensatedSum(f64) = .{};
for (advantages) |a| {
total.add(a);
}
const mean: f64 = total.value() / float64(n);
var spread: CompensatedSum(f64) = .{};
for (advantages) |a| {
spread.add((a - mean) * (a - mean));
}
// POPULATION deviation, not sample: this is the whole batch, not a sample from it, and
// the distinction is visible for a small minibatch.
const deviation: f64 = @sqrt(spread.value() / float64(n));
for (advantages) |*a| {
a.* = (a.* - mean) / (deviation + config.advantage_epsilon);
}
}
var objective: CompensatedSum(f64) = .{};
var kl: CompensatedSum(f64) = .{};
var entropy_total: CompensatedSum(f64) = .{};
var clipped: usize = 0;
for (0..n) |i| {
const r: f64 = log_prob_new[i] - log_prob_old[i];
const ratio: f64 = @exp(r);
const unclipped: f64 = ratio * advantages[i];
const bounded: f64 = zm.clamp(ratio, 1 - config.clip, 1 + config.clip) * advantages[i];
// THE MINIMUM, which is what makes it pessimistic: the objective takes the worse of the
// two branches, so a ratio that moved favourably stops paying once it leaves the trust
// region while an unfavourable one keeps counting.
objective.add(@min(unclipped, bounded));
if (ratio < 1 - config.clip or ratio > 1 + config.clip) {
clipped += 1;
}
kl.add(@exp(r) - 1 - r);
entropy_total.add(entropies[i]);
}
const count: f64 = @floatFromInt(n);
return .{
// NEGATED: the objective is maximised, a loss is minimised.
.policy_loss = -objective.value() / count,
.clip_fraction = float64(clipped) / count,
.approx_kl = kl.value() / count,
.entropy = entropy_total.value() / count,
};
}
The objective takes the worse of the clipped and unclipped branches. What that buys:
ratio 2.0, advantage +1 -> min(2.0, 1.2) = 1.2 improvement CAPPED
ratio 2.0, advantage -1 -> min(-2.0, -1.2) = -2.0 penalty NOT cappedA step that moved favourably stops paying once it leaves the trust region; an unfavourable one keeps counting in full. That asymmetry is what makes PPO safe — it will not let a good-looking update run away, but it will always punish a bad one.
A broken PPO run and a healthy one look the same from the loss, because the clipped objective is supposed to stop improving once the policy has moved far enough.
| clip_fraction | What it means |
|---|---|
| near 0 | the policy barely moved — the epochs after the first did nothing |
| near 1 | slamming the clip on every sample — the step is too large and the clip is all that holds it together |
approx_kl rising sharply is the usual trigger for PPO's early exit, and watching it
is cheaper than discovering a collapsed policy several thousand steps later.
The cheap one is mean(log_old - log_new). It is unbiased and it
can come out negative — and a negative divergence is not a small number, it
is a wrong one. A threshold on it either never fires or fires on noise.
r = log_new - log_old
cheap: mean(-r) can be negative
k3: exp(r) - 1 - r non-negative for every rexp(r) ≥ 1 + r holds everywhere, so the k3 estimator is non-negative by
construction rather than by luck. It is positive whether the policy moved up or down, which the
cheap one is not.
Raw advantages carry the scale of the reward. An environment paying 1 per step and one paying 1000 produce gradients a thousand times apart from the same policy error, so one learning rate cannot serve both.
Subtracting the mean and dividing by the deviation makes the step size a property of the algorithm rather than of the reward units. The deviation is the population one, not the sample one — this is the whole batch, not a sample from it, and the difference is visible for a small minibatch.
const y = try zn.dqnTarget(f64, transition, q_next_target, q_next_online, 0.99, .double);dqnTarget/// The regression target for one transition: `r + gamma * (1 - terminal) * Q(s', a')`.
///
/// WHY `double` EXISTS, AND WHAT IT FIXES
///
/// Vanilla DQN takes `max_a' Q_target(s', a')`. The maximum of noisy estimates is **biased
/// upward** - if every Q is right on average but wrong individually, the largest one is likely
/// to be one that happened to be over-estimated. Bootstrapping off it feeds that error back in,
/// and the values climb away from anything the environment pays.
///
/// Double DQN splits the two jobs: the ONLINE network picks which action looks best, the TARGET
/// network says what it is worth. The two networks have independent errors, so the selection
/// no longer systematically picks the target's over-estimates.
///
/// It is the same fix in one line, and it is why `double` is the sensible default for new code
/// even though `vanilla` is what the original paper did.
///
/// `q_next_online` is ignored for `vanilla` - pass it anyway so a caller switching between the
/// two changes one argument rather than restructuring.
pub fn dqnTarget(
comptime T: type,
transition: Transition,
q_next_target: []const T,
q_next_online: []const T,
gamma: T,
kind: DqnKind,
) Error!T {
comptime requireFloat(T);
if (q_next_target.len == 0 or q_next_online.len != q_next_target.len) {
return Error.ShapeMismatch;
}
if (gamma < 0 or gamma > 1) {
return Error.DomainError;
}
// A TERMINAL TRANSITION HAS NO FUTURE, so the target is the reward alone. Not
// `reward + gamma * something_small` - there is no next state, and its value is zero by
// definition rather than by estimate.
if (transition.terminal) {
return transition.reward;
}
const bootstrap: T = switch (kind) {
.vanilla => blk: {
var best: T = q_next_target[0];
for (q_next_target[1..]) |q| {
best = @max(best, q);
}
break :blk best;
},
.double => blk: {
// The ONLINE net chooses...
var chosen: usize = 0;
var best: T = q_next_online[0];
for (q_next_online[1..], 1..) |q, i| {
if (q > best) {
best = q;
chosen = i;
}
}
// ...and the TARGET net values that choice.
break :blk q_next_target[chosen];
},
};
return transition.reward + gamma * bootstrap;
}
vanilla: y = r + gamma*(1-terminal) * max_a' Q_target(s', a')
double: a* = argmax_a' Q_online(s', a')
y = r + gamma*(1-terminal) * Q_target(s', a*)Vanilla DQN takes the maximum of the target network's Q values. The maximum of noisy estimates is biased upward: if every Q is right on average but wrong individually, the largest one is likely to be one that happened to be over-estimated. Bootstrapping off it feeds that error back in, and the values climb away from anything the environment pays.
Double DQN splits the two jobs. The online network picks which action looks best; the target network says what it is worth. Their errors are independent, so the selection no longer systematically picks the target's over-estimates.
Take a state where the true value of every action is zero and give two networks independent noise around it. Over 400 trials with eight actions:
vanilla mean target > 1.0 the max of eight standard normals averages about 1.4
double mean target < 0.3 close to the truthWhen the two networks agree the estimators give the same answer, so switching is safe on a converged agent and differs only while the estimates are still noisy.
The target is the reward alone — not reward + gamma * something_small.
The next state does not exist, and its value is zero by definition rather than by estimate.
A target Q of 9 sitting in the array is ignored entirely, which is what a missing terminal flag
would let through.
A falling value loss says the critic is fitting something. It does not say the something is the returns.
critic predicts the mean return for every state
value loss small and shrinking
explained var exactly 0 it has learned no relationship at allNegative means worse than that constant baseline — a real and common state early in training, and invisible in the loss curve. Constant returns give zero rather than a division by zero: there was nothing to explain.
const y = try zn.sacTarget(f64, transition, next_estimates, next_log_prob,
gamma, alpha, .minimum);sacTarget/// The SAC regression target for one transition.
///
/// WHAT MAKES IT SAC RATHER THAN TD3
///
/// The entropy term. TD3's target is `r + gamma * min_i Q_i(s', a')`; SAC subtracts
/// `alpha * log p(a'|s')` from the bootstrap, so an action the policy was CONFIDENT about is
/// worth less than an equally-good action it was uncertain about.
///
/// That single term is what makes the objective "maximise return **and** keep exploring", and it
/// is why SAC does not need an exploration schedule bolted on: the incentive is in the value
/// function.
///
/// `next_estimates` is a slice for the reason `CriticAggregate` documents - two networks, M
/// dropout passes, or a random subset of an ensemble all arrive the same way.
pub fn sacTarget(
comptime T: type,
transition: Transition,
next_estimates: []const T,
next_log_prob: T,
gamma: T,
alpha: T,
how: CriticAggregate,
) Error!T {
comptime requireFloat(T);
if (gamma < 0 or gamma > 1 or alpha < 0) {
return Error.DomainError;
}
if (transition.terminal) {
// No next state, so neither its value NOR its entropy enters. Keeping the entropy term
// here would pay a bonus for the uncertainty of an action that was never taken.
return transition.reward;
}
const q: T = try aggregateCritics(T, next_estimates, how);
return transition.reward + gamma * (q - alpha * next_log_prob);
}
Every modern off-policy method attacks the same problem — Q values drifting upward, which caps how many gradient steps you can take per environment step. They differ only in where the diversity of estimates comes from:
| Method | Estimates come from | Combined by |
|---|---|---|
| TD3, SAC | two separate networks | minimum |
| DroQ | M dropout passes of one net | minimum |
| REDQ | a random M of an ensemble of N | minimum |
| CrossQ | one LayerNorm net, no target network | single value |
next_estimates is a slice, not a pair, so all four are the same
call. A single estimate passes through under either rule, which is what makes CrossQ a one-element
list rather than a special case.
A minimum of unbiased estimates under-estimates. Measured over noise around a true value of zero, with two critics:
minimum mean target < -0.3 the min of two standard normals averages -0.56
mean mean target ~ 0 unbiasedThe asymmetry is deliberate: an over-estimate gets bootstrapped into the next target
and compounds, while an under-estimate merely slows learning. .mean is there
for critics whose regularisation is elsewhere — a dropout or LayerNorm critic is already
conservative and does not need a second pessimism on top.
TD3: y = r + gamma * min_i Q_i(s', a')
SAC: y = r + gamma * (min_i Q_i(s', a') - alpha * log p(a'|s'))An action the policy was confident about is worth less than an equally-good action it
was uncertain about. That single term is why SAC needs no exploration schedule bolted on: the
incentive lives in the value function. Setting alpha to zero gives exactly TD3's
target, so the two methods differ by one coefficient.
A terminal transition drops the entropy term as well as the value — keeping it would pay a bonus for the uncertainty of an action that was never taken.
Using all N critics would be maximally pessimistic and slow to learn. A fixed M would let the same critics' shared error dominate every target. Re-drawing each step is what stops the pessimism being systematic — M = 2 of N = 10 is the paper's setting.
alpha is learned rather than set, because a fixed value trades return against
exploration at a rate that has to suit the reward scale, the action dimension and the stage
of training — and it suits at most one of those at a time.
The gradient is alpha * (entropy - target_entropy), and the usual target is
-action_dim. So with target_entropy = -2:
log_prob -5 entropy 5 above target -> gradient +, alpha falls
log_prob -1 entropy 1 still above -> gradient +, alpha falls less steeply
log_prob +3 entropy -3 below target -> gradient -, alpha rises
log_prob +2 entropy -2 AT the target -> gradient 0The fixed point is log_prob = +2, not −2. A continuous density can exceed 1,
so a positive log-probability is an ordinary state for a concentrated policy and is exactly what
-action_dim asks for.
log_alpha is the parameter, not alpha, for the same reason
DiagGaussian takes log_std: exp keeps it positive for free,
where a clamp would give a zero gradient exactly when it binds.
A Graph records operations as it performs them and walks the record backwards to
produce gradients. The same XOR network as the previous section, differentiated automatically:
var graph: zn.Graph(f64) = .init(scope.allocator());
const gx = try graph.constant(x);
const gt = try graph.constant(target);
const gw1 = try graph.parameter(w1);
const gb1 = try graph.parameter(b1);
const gw2 = try graph.parameter(w2);
const gb2 = try graph.parameter(b2);
const pre1 = try graph.add(try graph.matmul(gx, gw1), gb1);
const h = try graph.tanh(pre1);
const y = try graph.add(try graph.matmul(h, gw2), gb2);
const loss = try graph.mseLoss(y, gt);
try graph.backward(loss);
const dw1 = try graph.gradOf(gw1);A Var is a u32 index into the graph, not a pointer. Keying gradients
off the tensor would not work: a view shares data.ptr with its parent, so w
and w.transpose(0,1) would collide, and a Tensor is copied by value
throughout the library so it has no stable address at all. The graph owns the correspondence
between a value, its tensor and its gradient, and Tensor needs no change to support
it.
parameter allocates a gradient buffer; constant does not.
gradOf on a constant is DomainError rather than a buffer of zeros that
looks like an answer.
Each operation computes its result immediately and appends a node describing how to push a gradient back through it. Computing eagerly means every shape is known when it is needed, so outputs are allocated as the graph is built and there is no separate shape-inference pass to disagree with the forward pass.
Nodes are appended in execution order, and a value can only depend on values recorded before it. Walking the list backwards therefore visits every consumer of a value before the value itself — a topological order by construction, with no sort, no visited set and no cycle check.
backward takes a value of a single element and seeds its gradient with 1.
Differentiating a tensor with respect to parameters gives a Jacobian rather than a gradient;
reverse mode answers how one number changes, so the entry point requires one. Anything else is
UnsupportedShape. The pass zeroes every gradient before accumulating, so calling it
twice gives the same answer rather than doubling.
A bias of shape (1, n) added to (m, n) is used m times in the forward pass, so its gradient is the sum of m contributions. Accumulation reduces along any axis that was broadcast:
if (dest.size() == contribution.size()) {
return zimrnum.add(T, dest, dest, contribution);
}
// Broadcast along the leading axis: sum the contribution down to the bias's width.Adding the raw gradient instead would be low by a factor of m, and the network would still train — with a bias moving m times too slowly. The test asserts the bias gradients specifically for that reason.
| Operation | Backward |
|---|---|
| add | Both inputs receive the gradient; a broadcast input receives its sum. |
| sub | The left input receives the gradient, the right its negation. |
| mul | Each side receives the other side's value times the incoming gradient. |
| matmul | dA = g @ Bᵀ, dB = Aᵀ @ g |
| tanh | tanhGrad(out, g), using the forward output |
| sigmoid | sigmoidGrad(out, g), using the forward output |
| relu | reluGrad(in, g), using the forward input |
| mseLoss | 2(a − t)/n, scaled by the incoming gradient |
| crossEntropy | softmax − onehot, over the row count |
| softmaxRows | y · (g − ∑ g·y) per row; the Jacobian applied without being formed |
| layerNormRows | (1/σ) · (g − mean(g) − y · mean(g·y)) per row |
| scale | The gradient times the factor |
| dropout | The gradient times the same mask as the forward pass |
| transpose | The transpose of the gradient. Materialised, because a value of the graph must own its storage. |
The layer norm rule is checked against a finite difference on rows given very different scales and offsets, where a rule that dropped the variance term would be visibly wrong. Measured disagreement per operation, each alone: layer norm 9.7e-12, softmax 1.1e-11, scale 1.1e-9, dropout 7.1e-9.
dropout draws its mask once and stores it as a value of the graph. A dropout that
redrew on every recompute would be a different function each time, and
checkGradient would compare a derivative against a finite difference of something
else. resampleDropout draws fresh masks, which a training loop calls once per step.
The scaling is inverted — kept elements are multiplied by 1/(1 − rate)
— so removing the node at inference is exactly equivalent to rate = 0.
try graph.recompute();
const worst = try graph.checkGradient(loss, gw1, 1.0e-5);recompute replays the recorded operations forward, writing new results into the
buffers already allocated. Leaves are skipped, so writing a new batch into an input tensor and
calling recompute evaluates the same model on new data. Without it a graph would be
single-use and a training loop would rebuild the whole tape every step.
checkGradient perturbs each element of a parameter, replays, and compares a
central finite difference of the loss against the tape's gradient, returning the largest
disagreement. It is the one check that does not repeat the derivation: a hand-written backward
pass and a tape rule can be wrong in the same way, because the same person derived both, but a
finite difference does not know the derivative at all.
Central rather than forward: (f(x+h) − f(x−h)) / 2h has error
O(h²) where (f(x+h) − f(x)) / h has O(h). At h = 1e-5 in
f64 that is the difference between agreeing to ten digits and to three, which decides
whether a real error is visible above the noise. Every parameter of a graph exercising all eight
operations agrees to better than 1e-8.
The routine restores every value and recomputes before returning, so checking a model leaves it exactly as it was found.
const l1: zn.Conv2d(f64) = try .init(a, rng, .{ 3, 3 }, .{ 1, 1 }, .{ 1, 1 }, .he);
const h1 = try l1.attach(&graph, input); // .out, .kernel, .bias
const table: zn.Embedding(f64) = try .init(a, rng, vocabulary, width);
const e = try table.attach(&graph, &indices); // .out, .tableConv2d/// A 2-D convolution layer: a learned kernel and a scalar bias, on `Dense`'s template.
///
/// SAME SHAPE AS `Dense`: own the parameters, `attach` to a graph, hand back the handles the
/// optimiser needs. That is the whole contract of a layer here, and it is why there is no
/// `Sequential` - a struct holding two of these IS the composition, and `attach` chains.
///
/// Single channel, rank 2. The kernel is initialised by fan-in, counting the window: a 3x3
/// kernel sees nine inputs per output, so `He` over 9 is the right scale and `He` over the whole
/// image is not.
pub fn Conv2d(comptime T: type) type {
comptime requireFloat(T);
return struct {
const Self = @This();
kernel: Tensor(T),
bias: Tensor(T),
stride: [2]usize,
pad: [2]usize,
pub const Attached = struct {
out: Var,
kernel: Var,
bias: Var,
};
pub fn init(
gpa: Allocator,
rng: Rng,
size: [2]usize,
stride: [2]usize,
pad: [2]usize,
scheme: InitScheme,
) Error!Self {
const kernel: Tensor(T) = try Tensor(T).alloc(gpa, &size);
const bias: Tensor(T) = try Tensor(T).alloc(gpa, &.{ 1, 1 });
const fan_in: usize = size[0] * size[1];
switch (scheme) {
.xavier => try initXavier(T, kernel, rng, fan_in, fan_in),
.he => try initHe(T, kernel, rng, fan_in),
}
bias.fill(0);
return .{ .kernel = kernel, .bias = bias, .stride = stride, .pad = pad };
}
/// Convolve, then add the bias - which broadcasts from `(1, 1)` over the whole output.
pub fn attach(self: Self, graph: *Graph(T), input: Var) Error!Attached {
const k: Var = try graph.parameter(self.kernel);
const b: Var = try graph.parameter(self.bias);
const convolved: Var = try graph.conv2d(input, k, self.stride, self.pad);
return .{ .out = try graph.add(convolved, b), .kernel = k, .bias = b };
}
};
}
Embedding/// A lookup table of learned row vectors, indexed by whole numbers.
pub fn Embedding(comptime T: type) type {
comptime requireFloat(T);
return struct {
const Self = @This();
table: Tensor(T),
pub const Attached = struct {
out: Var,
table: Var,
};
/// Rows are drawn from the unit normal rather than by fan-in: an embedding row is not
/// a projection of anything, so there is no fan to scale by. The usual choice.
pub fn init(gpa: Allocator, rng: Rng, count: usize, width: usize) Error!Self {
const table: Tensor(T) = try Tensor(T).alloc(gpa, &.{ count, width });
rng.fillNormal(T, table.data);
return .{ .table = table };
}
pub fn attach(self: Self, graph: *Graph(T), rows: []const usize) Error!Attached {
const t: Var = try graph.parameter(self.table);
return .{ .out = try graph.embedding(t, rows), .table = t };
}
};
}
The backward pass is the forward loop with the accumulation reversed. Forward,
each output reads a window: y[i][j] += x[i·s+a−p][j·s+b−p] ·
k[a][b]. Both gradients fall out of that one line by moving what is read and what is
written:
dx[i·s+a−p][j·s+b−p] += g[i][j] · k[a][b]
dk[a][b] += g[i][j] · x[i·s+a−p][j·s+b−p]So it is written as one pass producing both, rather than as the "full correlation with a flipped kernel" the textbooks give. That form is correct and is a second thing to get right; this one cannot disagree with the forward pass about which taps are in range, because it asks the identical question.
An embedding's gradient is a scatter-add, and the += is the whole
point: a token appearing three times in a batch must accumulate three contributions into one row.
A scatter that assigned would train on the last occurrence and silently discard the rest. The test
indexes row 2 three times and lets the finite difference — which perturbs the row and sees
all three uses move — catch it.
const q = try graph.matmul(gx, gq);
const k = try graph.matmul(gx, gk);
const v = try graph.matmul(gx, gv);
const scores = try graph.scale(
try graph.matmul(q, try graph.transpose(k)),
1.0 / @sqrt(@as(f64, @floatFromInt(width))),
);
const weights = try graph.softmaxRows(scores);
const attended = try graph.matmul(weights, v);softmax(Q Kᵀ / √d) V is four operations the tape already has. Composing
them is what shows the tape composes: a gradient error in any one, or in how gradients
accumulate through a value used more than once, appears here. gx feeds all three
projections, so its gradient is the sum of three paths — the accumulation a single-use test
never exercises. All four parameters agree with a finite difference to better than 1e-8, and no
backward rule was written for attention itself.
try zn.initXavier(f64, w, rng, fan_in, fan_out); // uniform ±√(6/(in+out))
try zn.initHe(f64, w, rng, fan_in); // normal, deviation √(2/in)
const layer: zn.Dense(f64) = try .init(a, rng, 2, 16, .he);
const attached = try layer.attach(&graph, input);
// attached.out, attached.weight, attached.biasinitXavier/// Fill `w` with Glorot/Xavier uniform values: `+/-sqrt(6 / (fan_in + fan_out))`.
///
/// WHY THE SCALE DEPENDS ON THE SHAPE
///
/// A layer's output variance is roughly `fan_in` times its weight variance. Fill every layer from
/// the same distribution and that factor compounds: after ten layers the activations are either
/// saturated or vanishing, and the network cannot train regardless of the optimiser. Scaling by
/// the fan sizes holds the variance constant from layer to layer.
///
/// Xavier balances the FORWARD and BACKWARD passes, which is why both fans appear - the
/// forward pass wants `1/fan_in` and the backward wants `1/fan_out`, and `2/(fan_in + fan_out)`
/// is their harmonic compromise. It suits `tanh` and `sigmoid`, whose useful range is symmetric
/// about zero.
pub fn initXavier(
comptime T: type,
w: Tensor(T),
rng: Rng,
fan_in: usize,
fan_out: usize,
) Error!void {
comptime requireFloat(T);
if (fan_in == 0 or fan_out == 0) {
return Error.DomainError;
}
const total: T = @floatFromInt(fan_in + fan_out);
const limit: T = @sqrt(@as(T, 6.0) / total);
var i: usize = 0;
while (i < w.size()) : (i += 1) {
w.data[w.base + i] = rng.uniform(T, @intCast(i), -limit, limit);
}
}
initHe/// Fill `w` with He normal values: standard deviation `sqrt(2 / fan_in)`.
///
/// THE 2 IS THERE BECAUSE `relu` DISCARDS HALF ITS INPUT. A rectifier zeroes the negative
/// half, halving the variance that reaches the next layer; doubling the weight variance restores
/// it. Using Xavier with `relu` loses a factor of two per layer, which is a factor of a thousand
/// over ten layers.
pub fn initHe(comptime T: type, w: Tensor(T), rng: Rng, fan_in: usize) Error!void {
comptime requireFloat(T);
if (fan_in == 0) {
return Error.DomainError;
}
const deviation: T = @sqrt(@as(T, 2.0) / @as(T, @floatFromInt(fan_in)));
var i: usize = 0;
while (i < w.size()) : (i += 1) {
w.data[w.base + i] = rng.normal(T, @intCast(i)) * deviation;
}
}
Densepub fn Dense(comptime T: type) type {
comptime requireFloat(T);
return struct {
const Self = @This();
weight: Tensor(T),
bias: Tensor(T),
/// Allocate and initialise. `He` for a layer feeding a rectifier, `Xavier` otherwise.
pub fn init(
gpa: Allocator,
rng: Rng,
fan_in: usize,
fan_out: usize,
scheme: InitScheme,
) Error!Self {
const weight: Tensor(T) = try Tensor(T).alloc(gpa, &.{ fan_in, fan_out });
const bias: Tensor(T) = try Tensor(T).alloc(gpa, &.{ 1, fan_out });
switch (scheme) {
.xavier => try initXavier(T, weight, rng, fan_in, fan_out),
.he => try initHe(T, weight, rng, fan_in),
}
// Biases start at zero. A random bias adds nothing a random weight does not already
// provide, and a non-zero one shifts every unit's activation before any data is seen.
bias.fill(0);
return .{ .weight = weight, .bias = bias };
}
/// What `attach` hands back: the output, and handles for the two parameters.
///
/// RETURNING ONLY THE OUTPUT IS NOT ENOUGH, which a first attempt at a training loop
/// showed immediately: the loop needs the parameter handles to read their gradients, and
/// with the layer holding no `Var` there is nowhere else to get them. Returning all three
/// keeps the layer graph-agnostic AND gives the caller what a step requires.
pub const Attached = struct {
out: Var,
weight: Var,
bias: Var,
};
/// Record `input @ weight + bias` onto `graph`.
pub fn attach(self: Self, graph: *Graph(T), input: Var) Error!Attached {
const w: Var = try graph.parameter(self.weight);
const b: Var = try graph.parameter(self.bias);
return .{
.out = try graph.add(try graph.matmul(input, w), b),
.weight = w,
.bias = b,
};
}
};
}
A layer's output variance is roughly fan_in times its weight variance. Filling
every layer from the same distribution compounds that factor once per layer, and after ten layers
the activations are saturated or vanishing regardless of the optimiser. Both schemes scale by the
fan sizes to hold the variance constant.
Xavier balances the forward and backward passes — the forward wants
1/fan_in, the backward 1/fan_out, and 2/(fan_in + fan_out)
is their compromise. It suits tanh and sigmoid, whose useful range is
symmetric about zero.
He doubles that variance because relu discards its negative half,
halving what reaches the next layer. Using Xavier with a rectifier loses a factor of two per
layer — a factor of a thousand over ten.
Dense owns its weights as ordinary tensors and holds no Var. A
Var is meaningful to one graph, so storing one would tie a layer to the first graph
it met; keeping the layer graph-agnostic is what lets the same weights back two paths.
attach therefore returns the parameter handles alongside the output, because a
training loop needs them to read gradients and there is nowhere else to get them. Biases start at
zero: a random bias adds nothing a random weight does not, and shifts every unit before any data
is seen.
var opt: zn.Optimizer(f64) = try .init(
a, &graph,
&.{ a1.weight, a1.bias, a2.weight, a2.bias },
.{ .adam = .{ .rate = 0.02 } }, // or .{ .sgd = .{ .rate = 0.1 } }
);
while (step < 1000) : (step += 1) {
try graph.recompute();
try graph.backward(loss);
try opt.step(&graph);
}Optimizer/// An optimiser bound to a set of a graph's parameters, owning whatever state the algorithm
/// needs per parameter.
///
/// THE STATE LIVES HERE, NOT WITH THE CALLER
///
/// `sgdMomentum` needs a velocity per parameter and `adamStep` needs two moments; without this,
/// a caller training four parameters with Adam allocates and threads eight extra tensors by
/// hand, and a mistake - one moment buffer shared between two parameters - trains silently
/// wrong. Binding the state to the parameter list at construction makes that mistake
/// unrepresentable.
///
/// ONE TYPE, THREE ALGORITHMS. The choice is a field, not a type parameter, so a training loop
/// can be written once and the algorithm changed at one call site. The state buffers for an
/// algorithm that does not use them are never allocated.
///
/// `step` reads gradients through the graph and writes updates into the graph's parameter
/// tensors, so the next `recompute` sees them with no further bookkeeping.
pub fn Optimizer(comptime T: type) type {
comptime requireFloat(T);
return struct {
const Self = @This();
pub const Kind = union(enum) {
sgd: struct { rate: T },
momentum: struct { rate: T, momentum: T },
adam: Adam,
};
kind: Kind,
params: []const Var,
/// Per parameter: one buffer for momentum, two for Adam, none for SGD.
state: []Tensor(T),
/// Adam's 1-based step counter.
steps: usize,
pub fn init(
gpa: Allocator,
graph: *const Graph(T),
params: []const Var,
kind: Kind,
) Error!Self {
const per: usize = switch (kind) {
.sgd => 0,
.momentum => 1,
.adam => 2,
};
// THE PARAMETER LIST IS COPIED, NOT RETAINED
//
// The first version stored the caller's slice. Every test passed, because a test
// writes `&.{ a, b, c }` in the same frame that calls `step`. The GPU training example
// built its optimiser in `init` and stepped it from `update`, and the array literal
// was gone - an out-of-bounds read inside `valueOf`, on the first frame. A type that
// outlives the call that made it owns what it needs.
const owned: []Var = try gpa.alloc(Var, params.len);
@memcpy(owned, params);
const state: []Tensor(T) = try gpa.alloc(Tensor(T), params.len * per);
for (params, 0..) |v, i| {
const shape: []const usize = graph.valueOf(v).shape[0..graph.valueOf(v).rank];
var k: usize = 0;
while (k < per) : (k += 1) {
state[i * per + k] = try Tensor(T).alloc(gpa, shape);
state[i * per + k].fill(0);
}
}
return .{ .kind = kind, .params = owned, .state = state, .steps = 0 };
}
/// Apply one update to every bound parameter, from the gradients `backward` left.
pub fn step(self: *Self, graph: *const Graph(T)) Error!void {
self.steps += 1;
for (self.params, 0..) |v, i| {
const weight: Tensor(T) = graph.valueOf(v);
const grad: Tensor(T) = try graph.gradOf(v);
switch (self.kind) {
.sgd => |h| try sgdStep(T, weight, weight, grad, h.rate),
.momentum => |h| try sgdMomentum(
T,
weight,
weight,
grad,
self.state[i],
h.rate,
h.momentum,
),
.adam => |h| try adamStep(
T,
weight,
weight,
grad,
self.state[i * 2],
self.state[i * 2 + 1],
h,
self.steps,
),
}
}
}
};
}
An Optimizer is bound to a list of a graph's parameters at construction and owns
whatever state the algorithm needs for each: nothing for SGD, a velocity for momentum, two
moments for Adam. Without it a caller training four parameters with Adam threads eight extra
tensors by hand, and a mistake — one moment buffer shared between two parameters —
trains silently wrong. Binding the state to the parameter list makes that unrepresentable.
The algorithm is a field rather than a type parameter, so a loop is written once and the
algorithm changed at one site. State for an algorithm that does not use it is never allocated.
step reads gradients through the graph and writes updates into the graph's own
parameter tensors, so the next recompute sees them.
SGD through the optimiser is asserted equal to sgdStep written by hand
— the wrapper adds no arithmetic, so a difference would be a bookkeeping error in which
parameter received which gradient. The three algorithms on the classifier of the previous
section, from identical weights:
| Steps | SGD | Momentum | Adam |
|---|---|---|---|
| 50 | 0.769 | 0.580 | 0.539 |
| 200 | 0.628 | 0.177 | 0.043 |
| 1000 | 0.354 | 0.0089 | 0.0012 |
The ordering holds at 200 steps. Adam's two buffers per parameter are a cost; the table is what they buy.
The graph is built once, outside the loop. Each step replays it, differentiates it, and steps the parameters in place:
while (step < 3000) : (step += 1) {
try graph.recompute();
try graph.backward(loss);
try zn.sgdStep(f64, w1, w1, try graph.gradOf(gw1), 0.2);
try zn.sgdStep(f64, b1, b1, try graph.gradOf(gb1), 0.2);
try zn.sgdStep(f64, w2, w2, try graph.gradOf(gw2), 0.2);
try zn.sgdStep(f64, b2, b2, try graph.gradOf(gb2), 0.2);
}Nothing is allocated after the first iteration. sgdStep writes into the parameter
tensors the graph already holds, so the next recompute sees the updated weights
without any bookkeeping.
A three-class classifier on 32 points, labelled by quadrant sign, trained this way:
| Step | Loss | Accuracy |
|---|---|---|
| 0 | 1.282 | 0.22 |
| 200 | 0.126 | 0.97 |
| 3000 | 0.021 | 1.00 |
The gradient check runs before and after: 3.8e-11 at the start and 9.7e-12 after training, at a point far from where it began. A model that trains is not by itself evidence that its gradients are right — a wrong but correlated gradient still reduces a loss — so the two are checked separately.
const top = try graph.slice(x, 0, 0, 2); // axis, start, len
const bottom = try graph.slice(x, 0, 2, 3);
const back = try graph.concat(top, bottom, 0);Tensor.slice has always given a strided view for free. What it could not do is carry
a gradient — so anything slicing a computed value was out of reach.
That absence shaped two designs in this file: LstmCell could not fuse
its four gates into one matmul, and Attention could not split a projection into heads.
Both said "stated rather than half-built" and moved on. znum has neither op either.
The gradient lands in the rows the slice read, and every other row gets exactly
zero — a row nobody read influenced nothing. concat's backward is the
same thing read the other way: each input takes back the part of the gradient that came from its own
rows.
A scatter that writes to the wrong offset still runs, still produces numbers of the right shape, and trains something adjacent to what you meant. So the comparison is against a central difference at every element, and separately asserts the zeros: slice the middle two rows of five, and rows 0, 3 and 4 must have no gradient — not small, zero.
It also checks that slicing and rejoining is the identity on values. That matters because it means a wrong offset in either backward shows up in the gradient while the forward pass still looks perfect.
const h = try graph.silu(x); // x * sigmoid(x)
const h = try graph.gelu(x); // the tanh formsilu, gelu, softplus and elu all existed at
tensor level with no tape node at all. A network using any of them could be
evaluated and not trained — the forward works, the loss falls for the layers that do
have nodes, and the activation simply never learns. Nothing reports it.
That is a worse failure than a missing gradient function, because it looks like training.
Two of the four are compositions of nodes that already exist, so they cost nothing and their gradients are the tape's.
exp and log are on the tape now. They unlock compositions it could not
express — mish is x * tanh(softplus(x)), and a ratio of exponentials
becomes ordinary arithmetic.
But having the primitives does not mean composing them is safe. Written as
log(1 + exp(x)), softplus overflows: measured at f32, exp(89) is infinity,
so the composition returns inf from x = 89 upward where the true answer is
89 — softplus is very nearly the identity for large x.
So softplus is its own node. The identity behind its form:
log(1 + e^x) = max(x, 0) + log(1 + e^-|x|)
for x > 0: log(1 + e^x) = log(e^x (e^-x + 1)) = x + log(1 + e^-x)
for x ≤ 0: max is 0 and |x| = -x, leaving the originalThe exponential is therefore only ever taken of a negative number and cannot overflow.
But overflow is not the only failure. At x = −50 the answer is
1.93e−22, and any form computing log(1 + e^-|x|) returns exactly zero
— 1 + 1.9e-22 rounds to 1 in f64 and its logarithm is 0. znum's softplus does
this; so did ours, with a threshold at 20 that also left a 1e−10 discontinuity just above
it.
log1p takes the small argument directly and never forms the 1:
x log1p form log(1 + ...) form
-50 1.929e-22 0
-700 9.860e-305 0
-745 4.941e-324 0The tail survives to the smallest subnormal. It is the same trick logSumExp uses
one screen away and the one log(softmax(x)) needed — three places in this
library where forming the 1 loses everything.
The backward is what makes a dedicated node cheap: d/dx softplus(x) is exactly
sigmoid(x) — no branch, nothing to derive.
The test probes all four gradients separately, because the three backwards read three different
things: exp uses its own output, log needs the input, and
softplus uses sigmoid of the input. Picking the wrong one is easy and
gives plausible numbers. It also checks log(exp(x)) == x, which would fail if either
node were wrong in a way the separate probes happened to miss.
elu is a real node, not a composition: it switches on the sign, and there
is no elementwise select to build that from. Its backward is 1 above zero
and alpha * exp(x) below, read from the input — the output would
give the same numbers, but only the input says which side of the branch a value was on.
The shared gradient probe above uses all-positive values, because log needs them.
On positive input ELU is the identity, so that probe exercises one side of the
branch and reports success. It gets its own spread straddling zero —
{-3, -1, -0.1, 0, 0.1, 1, 3} — which is the only place the branch shows.
Two properties pin it beyond the derivative: above zero the output equals the input
exactly, and below zero it saturates toward -alpha without reaching it.
Every activation in the library now has a tape node —
relu, sigmoid, tanh, silu, gelu,
softplus, mish, elu. Before this pair of turns, four of those
could be evaluated and never trained.
GELU is written x * sigmoid(2u) rather than 0.5x * (1 + tanh(u)).
They are the same function:
(1 + tanh u)/2 = (1 + (e^u - e^-u)/(e^u + e^-u))/2
= (2e^u / (e^u + e^-u))/2
= 1 / (1 + e^-2u)
= sigmoid(2u)The first form needs a tensor of ones to add — allocated on every call, and on a graph rebuilt each training step that is a per-step allocation for a constant. The second needs none and is two nodes shorter.
GELU is the tanh form, not the erf definition. The exact version needs
erf; the approximation is what PyTorch, JAX and every published transformer actually
run, so matching it matters more than matching the definition — weights were trained against
it.
One might expect the tensor gelu was the erf definition and expected the two to differ
by about 1e−3. They agree to 8e−13 — zimrnum's tensor version is
the same approximation. The bar is set from that measurement rather than from an assumption about the
code.
And the finite-difference probe mutates its input, so the input must be
allocated. Every other test here builds tensors with
@constCast(&[_]f64{...}), which is safe because none of them writes — a
literal lives in read-only memory and @constCast only silences the type system.
Writing through it aborts.
Every piece above has its own test. None of those catches a mistake in how the pieces fit — a gradient with the wrong sign, a transpose on the wrong operand, a bias summed along the wrong axis. Each of those leaves every individual function correct and the network unable to learn.
The following is a test in src/zimrnum.zig. It trains a two-layer network on XOR,
which is the smallest problem a linear model cannot solve: without a working hidden layer the loss
stalls near 0.25.
// forward
try zn.matmul(f64, pre1, x, w1);
try zn.add(f64, pre1, pre1, try b1.broadcastTo(&.{ samples, hidden }));
try zn.tanh(f64, h, pre1);
try zn.matmul(f64, y, h, w2);
try zn.add(f64, y, y, try b2.broadcastTo(&.{ samples, 1 }));
const loss = try zn.mseLoss(f64, y, target);
// backward: d(mse)/dy is 2(y - t)/n
try zn.sub(f64, dy, y, target);
try zn.scale(f64, dy, dy, 2.0 / @as(f64, @floatFromInt(samples)));
try zn.matmul(f64, dw2, try h.transpose(0, 1), dy);
try zn.sumAxis(f64, db2, dy, 0);
try zn.matmul(f64, dh, dy, try w2.transpose(0, 1));
try zn.tanhGrad(f64, dpre, h, dh);
try zn.matmul(f64, dw1, try x.transpose(0, 1), dpre);
try zn.sumAxis(f64, db1, dpre, 0);
// update
try zn.sgdStep(f64, w1, w1, dw1, rate);
try zn.sgdStep(f64, b1, b1, db1, rate);
try zn.sgdStep(f64, w2, w2, dw2, rate);
try zn.sgdStep(f64, b2, b2, db2, rate);Weights and intermediates come from an arena, so nothing in the loop is freed individually and the whole allocation goes when the scope closes. The transposes are stride swaps — no copy.
| Step | Loss |
|---|---|
| 0 | 0.541 |
| 100 | 6.2e-6 |
| 500 | 4e-8 |
| 1000 | 0 |
Final outputs are −0.000000, 1.000000, 1.000000, −0.000000 against targets of 0, 1, 1, 0, and the weight gradients have fallen to about 2e-11.
One detail worth noting: tanhGrad takes h, the forward
output, not pre1. Passing the input compiles, runs, and produces a network
that learns far too slowly — the kind of error only an end-to-end check finds.
Tasks rather than functions. Each is the shortest correct way to do something the library is for, with a link to the section that explains the pieces.
| Task | How |
|---|---|
| A tensor of zeros, ones, or a constant | zeros, ones, full — 4.6 |
| A range, or n evenly spaced points | arange (stop excluded), linspace (endpoint named) — 4.6 |
| Same shape as an existing tensor | zerosLike, onesLike, fullLike |
| Convert element type | astype — out of range is an error — 4.7 |
| Save or load a tensor for numpy | saveNpy, loadNpy |
| Save a trained network | saveParameters, loadParameters — 11.18 |
| Task | How |
|---|---|
| A column where some values are absent | Series; valid == null means all present — 12.4 |
| Build one from a tensor with NaNs | fromTensorMarkingNaN |
| Sum or average, skipping gaps | sumValid, meanValid |
| Carry the last known value forward | fillForward — a leading gap stays missing — 12.6 |
| Drop the rows with gaps | dropMissing |
| Skip NaNs in a plain tensor | nanSum, nanMean — 9.8 |
| Task | How |
|---|---|
| Group rows by a key and aggregate | groupBy then aggregate — the result keeps its type — 12.7 |
| Turn labels into integer codes | factorize — 12.3 |
| The eight summary numbers | describe — 12.8 |
| Match rows of two tables on a key | joinRows then takeRows — 12.9 |
| Filter, sort, sample, or take the first n | all one function: takeRows with an index list; firstRows for the head |
| Yesterday's value, or a smoothed series | shiftRows, rollingMean — edges are missing — 12.11 |
| Task | How |
|---|---|
| The smallest complete example | 3.1 |
| A feed-forward stack | Chain — parameters collected, not counted — 11.12 |
| Stop one bad batch ruining the run | clipGradNorm before the step — 11.19 |
| An optimiser rather than a hand-written update | adamWStep, rmspropStep — 11.7 |
| Decay the learning rate | cosineLearningRate, warmup — 11.10 |
| Check a gradient is right | central differences against gradOf — 17.2 |
| Classification scores | confusionMatrix then perClassMetrics — 12.10 |
zimrnum compiles to SPIR-V, so a kernel can call the same function the CPU does rather than
restating it. The pattern is three lines in build.zig and one in the kernel:
// build.zig: mark the kernel module as wanting zimrnum
.wants_zimrnum = true,
// the kernel: call the host function directly
pub fn cartpole_step(c: k.Ctx(@This())) void {
if (c.id >= c.params.count) return;
state[c.id] = zn.cartpoleStep(state[c.id], action[c.id], dt);
}cartpoleStep/// The cart-and-pole control benchmark: push a cart left or right to keep a pole upright.
///
/// A PURE FUNCTION - THE STATE GOES IN AND THE NEXT STATE COMES OUT
///
/// znum's step takes `state: *Tensor(f32)` and updates it IN PLACE, with a doc comment warning
/// that the tensor "holds the observation for step t only until this returns, so a rollout copies
/// it out BEFORE calling". **That is a footgun stated in prose**: the caller has to remember, and
/// forgetting produces a rollout whose observations are all one step late, which trains and
/// converges to something slightly wrong.
///
/// Returning the next state makes it unforgettable. It also makes the step replayable - the same
/// state and action give the same answer, always - which is what a test needs and what a
/// debugger needs.
///
/// RESET IS A SEPARATE FUNCTION, AND THAT REMOVES TWO FLAGS
///
/// znum folds the reset into the step, so the step needs a `force_reset` parameter, and then a
/// `no_transition` special case because a forced reset produces no reward and no termination.
/// Both disappear when starting an episode is its own call: `cartpoleReset` gives a state, and
/// `cartpoleStep` only ever steps.
///
/// RADIANS HERE, AND THE SUFFIX SAYS SO - BECAUSE THIS ANGLE IS DIFFERENTIATED
///
/// The rest of this codebase moved to turns, and for geometry that was right: a quarter turn is
/// exactly 0.25, and `sinTurns` is exact where `sin(tau/4)` is not.
///
/// **A differentiated angle is the other case.** `d/dtheta sin(theta)` is `cos(theta)` only in
/// radians; in turns it is `tau * cos`, and every derivative in a dynamics system picks up that
/// factor. The cartpole equations are differential equations, so radians are the unit in which
/// they have no constants in them at all.
///
/// The first version of this held the pole in turns and converted the angular acceleration from
/// radians on the way in. That conversion was a place to be wrong, and I was: **the pole fell tau
/// times too fast** and every sign still looked right. Choosing the unit the equations are
/// written in removes the conversion rather than fixing it.
///
/// So: **turns for angles that are drawn, radians for angles that are differentiated**, and the
/// name carries the unit either way. `pole_rad` cannot be mistaken for `pole_turns`.
///
/// The limit is still not a rounded constant. znum writes `0.2095`, which is 12.0035 degrees;
/// `radFromDeg(12)` is twelve degrees exactly as far as the type can express it.
///
/// The dynamics are the standard ones - Barto, Sutton and Anderson 1983 - and the integration is
/// SEMI-IMPLICIT Euler: the position update uses the NEW velocity. Gym's default is explicit
/// Euler, which drifts; znum uses semi-implicit too and is right to.
pub fn cartpoleStep(
comptime T: type,
state: CartpoleState(T),
push: Push,
) CartpoleStep(T) {
comptime requireFloat(T);
const gravity: T = 9.8;
const cart_mass: T = 1.0;
const pole_mass: T = 0.1;
const total_mass: T = cart_mass + pole_mass;
const half_pole: T = 0.5;
const pole_moment: T = pole_mass * half_pole;
const push_force: T = 10.0;
const seconds: T = 0.02;
const force: T = switch (push) {
.left => -push_force,
.right => push_force,
};
// `zm.` because zimrnum has its own `sinRad` and it is a tensor op - the disambiguation the
// two `lint:off` lines at the top of this file exist for.
const lean: T = zm.sinRad(state.pole_rad);
const upright: T = zm.cosRad(state.pole_rad);
// NO CONVERSION ANYWHERE, WHICH IS THE POINT OF CHOOSING RADIANS
//
// These are the textbook equations transcribed, and every quantity is already in the unit the
// derivation uses. An earlier version held the pole in turns and converted here; that line
// was a place to be wrong and it was.
const shared: T = (force + pole_moment * state.pole_rate_rad * state.pole_rate_rad * lean) /
total_mass;
const pole_accel: T = (gravity * lean - upright * shared) /
(half_pole * (4.0 / 3.0 - pole_mass * upright * upright / total_mass));
const cart_accel: T = shared - pole_moment * pole_accel * upright / total_mass;
// Semi-implicit: the position uses the UPDATED velocity, which conserves energy far better
// than the explicit form over a long episode.
const cart_rate: T = state.cart_rate + seconds * cart_accel;
const pole_rate_rad: T = state.pole_rate_rad + seconds * pole_accel;
const next: CartpoleState(T) = .{
.cart = state.cart + seconds * cart_rate,
.cart_rate = cart_rate,
.pole_rad = state.pole_rad + seconds * pole_rate_rad,
.pole_rate_rad = pole_rate_rad,
};
// Twelve degrees, converted rather than rounded - and 2.4 metres of track either side.
const fall_rad: T = zm.radFromDeg(@as(T, 12.0));
const track: T = 2.4;
const failed: bool = @abs(next.cart) > track or @abs(next.pole_rad) > fall_rad;
return .{
.state = next,
// A reward of one for every step survived, which is the standard scoring and makes the
// return equal to the episode length - so "did it learn" reads directly off the number.
.reward = 1.0,
.failed = failed,
};
}
This works when the host function has a pure scalar form — one input struct to one output struct. A whole-tensor operation expressed over shapes and strides is a different algorithm from a kernel indexed by thread id, and sharing code between them gains nothing.
Every operation with a kernel is checked against its CPU twin on the same input, with a per-operation tolerance. Operations whose arithmetic cannot drift carry a bar of exactly zero — 9.10.
Every entry here is something that produced a wrong answer without producing an error, during the writing of this library. They are collected because a trap you meet at three in the morning has to be findable, not discoverable by reading the section it happens to live in.
Each links to where it is explained.
isValid before reading
a value — 19.1 / 12.6defer gpa.free(col.valid.?) written before the fill
unwraps a null — 12.6all of an empty tensor is TRUE and any is false.
Forced by composition, not chosen — and getting it backwards passes every non-empty test
— 9.8arange excludes its stop; linspace asks. Same count,
different spacing — 4.6log(1 + small) is zero once
small falls under the epsilon of 1. It has cost three separate values here —
log(softmax), softplus, and what logSumExp avoided by
design. log1p and the max-subtraction are the fixes —
17.11exp(-x) reaches zero near x = 103 in f32, so a value 110 below a
row maximum softmaxes to exactly zero and its logarithm is -inf —
8.8const is folded at
full precision and proves nothing about the type you named@constCast on a literal gives read-only memory. Fine to read from,
aborts on a write — and a finite-difference probe writes —
17.11astype is an error, not a wrap. numpy wraps;
Zig's @intFromFloat is undefined out of range, so the choice was forced —
4.7stack grows the rank; concat grows an axis. Both take
a list and an axis and the shapes look alike — 7.11swapaxes returns a VIEW. Writing through it writes through to the
original — 7.11d/dθ sin θ is cos θ only in radians, and the name
carries the unit either wayThe complete public surface, in source order, generated by tools/zimrnum_ref.zig.
A test in src/zimrnum.zig compares this table against the source and fails the build
if either gains an entry the other lacks. Names appearing twice — init,
transpose — belong to different types; the signature column distinguishes
them.
The GPU column names the kernel or kernels compared against that function on a
device, derived from the sweep in examples/zimrnum_field: each of its rows names a
kernel and calls one zn function. host means the function
runs on the CPU only. For factorisations, statistics, views and the autograd machinery that is by
design; for an elementwise or row-wise operation it means a kernel has not been written yet.
| Declaration | Signature | GPU | Notes |
|---|---|---|---|
Error | Error = error | host | One error set for the whole library. |
isFloat | isFloat(comptime T: type) bool | host | True for f16, f32, f64. |
requireFloat | requireFloat(comptime T: type) void | host | Compile-time guard; fails at the call site with the offending type named. |
requireInt | requireInt(comptime T: type) void | host | Whole numbers only. znum's bitwiseAnd has no guard at all. |
requireNumeric | requireNumeric(comptime T: type) void | host | Integers welcome. Only where the widening is TESTED at both widths. |
Ddof | Ddof = enum | host | population or sample. |
divisor | divisor(self: Ddof, n: usize) ?f64 | method | Returns null where the divisor is undefined. |
approxEqAbs | approxEqAbs(comptime T: type, a: T, b: T, tol: T) bool | host | Absolute tolerance. NaN equals nothing. |
Rng | Rng = struct | host | Counter-based. A draw is a pure function of (seed, index). |
init | init(seed: u32) Rng | method | Rng.init(seed) and Ctx.init(gpa, seed). |
split | split(self: Rng, label: u32) Rng | method | A derived generator from a label. |
bits | bits(self: Rng, index: u32) u32 | method | Raw 32-bit draw, stream 0. |
unitFloat | unitFloat(self: Rng, comptime T: type, index: u32) T | method | [0, 1) with the full mantissa of T. |
uniform | uniform(self: Rng, comptime T: type, index: u32, low: T, high: T) T | method | [low, high). |
normal | normal(self: Rng, comptime T: type, index: u32) T | method | Box–Muller. Its own streams. |
intBelow | intBelow(self: Rng, index: u32, bound: u32) u32 | method | [0, bound), unbiased. 0 when bound is 0. |
fillUniform | fillUniform(self: Rng, comptime T: type, out: []T) void | method | Indices 0..out.len. |
fillNormal | fillNormal(self: Rng, comptime T: type, out: []T) void | method | Indices 0..out.len. |
Ctx | Ctx = struct | host | { gpa, rng } — what a call carries. |
init | init(gpa: Allocator, seed: u32) Ctx | method | Rng.init(seed) and Ctx.init(gpa, seed). |
derive | derive(self: Ctx, label: u32) Ctx | method | Same allocator, independent stream. |
max_rank | max_rank: usize = 6; | host | 6. |
broadcastShape | broadcastShape(a: []const usize, b: []const usize, out_shape: *[max_rank]usize) Error!u8 | host | The NumPy rule. Returns the rank. |
Tensor | Tensor(comptime T: type) type | host | A view of T-typed storage. Owns nothing. |
alloc | alloc(gpa: Allocator, shape: []const usize) Error!Self | method | Dense row-major. The caller frees .data. |
fromSlice | fromSlice(data: []T, shape: []const usize) Error!Self | method | Extents must multiply to exactly data.len. |
size | size(self: Self) usize | method | Total elements addressed. |
isContiguous | isContiguous(self: Self) bool | method | Ask before walking memory linearly. |
flatIndex | flatIndex(self: Self, indices: []const usize) Error!usize | method | base + Σ stride·index. |
at | at(self: Self, indices: []const usize) Error!T | method | Bounds-checked read. |
offsetOf | offsetOf(self: Self, indices: []const usize) usize | method | No error: for a coordinate whose rank and bounds you already know. |
at2 | at2(self: Self, i: usize, j: usize) T | method | Rank named, so NO error: the bounds check is Zig's, where Zig puts it. |
setAt2 | setAt2(self: Self, i: usize, j: usize, value: T) void | method | The mirror of at2. |
at1 | at1(self: Self, i: usize) T | method | Rank-1 read, error-free. |
setAt1 | setAt1(self: Self, i: usize, value: T) void | method | Rank-1 write, error-free. |
setAt | setAt(self: Self, indices: []const usize, value: T) Error!void | method | Bounds-checked write. |
fill | fill(self: Self, value: T) void | method | Handles strided views, not only dense ones. |
reshape | reshape(self: Self, shape: []const usize) Error!Self | method | Contiguous only; same element count. |
isAliased | isAliased(self: Self) bool | method | True when a stretched axis has extent > 1. Blocks writes. |
broadcastTo | broadcastTo(self: Self, shape: []const usize) Error!Self | method | Stretches with stride 0. Read-only in practice. |
slice | slice(self: Self, axis: usize, start: usize, len: usize) Error!Self | method | A window. Shares storage. |
permute | permute(self: Self, order: []const usize) Error!Self | method | Reorders axes. order must be a permutation. |
squeeze | squeeze(self: Self, axis: usize) Error!Self | method | Removes a size-1 axis. |
unsqueeze | unsqueeze(self: Self, axis: usize) Error!Self | method | Inserts a size-1 axis. |
flatten | flatten(self: Self) Error!Self | method | Rank-1 view. Contiguous only. |
moveAxis | moveAxis(self: Self, from: usize, to: usize) Error!Self | method | The other axes keep their order. |
transpose | transpose(self: Self, a: usize, b: usize) Error!Self | method | Exchanges two axes. |
map | map(comptime T: type, out: Tensor(T), a: Tensor(T), comptime f: fn (T) T) Error!void | transpose | Unary elementwise with a comptime function. |
zip | zip(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T), comptime f: fn (T, T) T) Error!void | host | Binary elementwise. Broadcasts both inputs to out. |
add | add(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | add, bcast_add | |
sub | sub(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | sub | |
mul | mul(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | mul | Elementwise, not a matrix product. |
div | div(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | div | Float only. |
CompensatedSum | CompensatedSum(comptime T: type) type | host | Neumaier accumulator. Every sum in the library goes through it. |
add | add(self: *@This(), x: T) void | method | |
value | value(self: @This()) T | method | total + compensation. |
sumAll | sumAll(comptime T: type, a: Tensor(T)) T | sum_all, sum_all_tiled | Compensated (Neumaier). |
sumAllFast | sumAllFast(comptime T: type, a: Tensor(T)) T | host | Pairwise. Faster, wrong under cancellation. |
meanAll | meanAll(comptime T: type, a: Tensor(T)) Error!T | mean_all | DomainError when empty. |
minAll | minAll(comptime T: type, a: Tensor(T)) Error!T | min_all | NaN is not ordered. |
maxAll | maxAll(comptime T: type, a: Tensor(T)) Error!T | max_all | NaN is not ordered. |
meanAxis | meanAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void | mean_axis0 | |
maxAxis | maxAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void | max_axis0 | NaN never wins. |
minAxis | minAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void | min_axis0 | |
prodAxis | prodAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void | host | |
varianceAxis | varianceAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize, ddof: Ddof) Error!void | variance_axis0 | Two-pass. ddof selects the divisor. |
prodAll | prodAll(comptime T: type, a: Tensor(T)) T | prod_all | |
cumsum | cumsum(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | cumsum_rows | Along the last axis. Rank 1 or 2. |
anyNonzero | anyNonzero(comptime T: type, a: Tensor(T)) bool | host | |
allNonzero | allNonzero(comptime T: type, a: Tensor(T)) bool | host | True for an empty tensor. |
where | where(comptime T: type, out: Tensor(T), mask: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | mask, a, b all broadcast to out. |
correlationMatrix | correlationMatrix(comptime T: type, out: Tensor(T), columns: []const Tensor(T)) Error!void | host | Every pair. The diagonal is ASSIGNED 1, not computed. |
coefficientOfVariation | coefficientOfVariation(comptime T: type, a: Tensor(T), ddof: Ddof) Error!T | host | Deviation as a fraction of the mean. A zero mean has no scale: DomainError. |
medianAbsDev | medianAbsDev(comptime T: type, scratch: []T, a: Tensor(T)) Error!T | host | The robust one. A separate name from meanAbsDev rather than a flag. |
standardError | standardError(comptime T: type, a: Tensor(T), ddof: Ddof) Error!T | host | Deviation shrunk by sqrt(N). |
mode | mode(comptime T: type, scratch: []T, a: Tensor(T)) Error!T | host | Most common value; ties go to the SMALLEST, reproducibly. |
Series | Series(comptime T: type) type | host | A named column. `valid == null` means all present - no mask, no cost. |
fromSlice | fromSlice(gpa: Allocator, name: []const u8, data: []const T) Error!Self | method | Extents must multiply to exactly data.len. |
fromTensorMarkingNaN | fromTensorMarkingNaN(gpa: Allocator, name: []const u8, source: Tensor(T)) Error!Self | method | NaN in, honest mask out. Allocates only if something IS missing. |
deinit | deinit(self: Self) void | method | Frees the values and the mask, if there is one. |
len | len(self: Self) usize | method | Entries, present or not. |
isValid | isValid(self: Self, i: usize) bool | method | No mask means everything is present. |
sumValid | sumValid(self: Self) struct | method | SKIPS the gaps. Raw sumAll adds zero for an integer gap and looks fine. |
meanValid | meanValid(self: Self) ?T | method | Null, not NaN, when nothing is present. |
fillMissing | fillMissing(self: *Self, value: T) void | method | A constant into every gap. The mask is then freed. |
fillForward | fillForward(self: *Self) void | method | Last real value forward. A LEADING gap stays missing - nothing to borrow. |
fillBackward | fillBackward(self: *Self) void | method | The mirror. A trailing gap stays missing. |
dropMissing | dropMissing(self: Self, gpa: Allocator, new_name: []const u8) Error!Self | method | A NEW shorter column: dropping rows changes the length. |
countValid | countValid(self: Self) usize | method | How many entries are really there. |
Column | Column = union(enum) | host | A tagged union. `inline else` stamps one arm per type, not six copies. |
deinit | deinit(self: Column) void | method | Frees the values and the mask, if there is one. |
len | len(self: Column) usize | method | Entries, present or not. |
name | name(self: Column) []const u8 | method | The column name, whatever its element type. |
countValid | countValid(self: Column) usize | method | How many entries are really there. |
Frame | Frame = struct | host | Named columns, all one height. A ragged table is refused at the door. |
init | init(gpa: Allocator) Frame | method | Rng.init(seed) and Ctx.init(gpa, seed). |
deinit | deinit(self: *Frame) void | method | Frees the values and the mask, if there is one. |
width | width(self: *const Frame) usize | method | How many columns. |
height | height(self: *const Frame) usize | method | How many rows - every column agrees. |
addColumn | addColumn(self: *Frame, taken: Column) Error!void | method | Refuses a wrong height AND a duplicate name: a lookup must not be ambiguous. |
find | find(self: *const Frame, wanted: []const u8) ?usize | method | Index of a named column, or null. |
rename | rename(self: *Frame, from: []const u8, to: []const u8) Error!void | method | Refuses a collision; renaming to itself is a no-op. |
dropColumn | dropColumn(self: *Frame, name: []const u8) Error!void | method | Order preserved - swapping the last into the hole would reorder silently. |
columnNamed | columnNamed(self: *const Frame, wanted: []const u8) ?Column | method | The column itself, or null. |
GroupBy | GroupBy(comptime K: type) type | host | order + starts with a TRAILING SENTINEL, so the last group is not special. |
deinit | deinit(self: Self) void | method | Frees the values and the mask, if there is one. |
groupCount | groupCount(self: Self) usize | method | starts.len - 1, thanks to the sentinel. |
rows | rows(self: Self, g: usize) []const usize | method | The row indices in one group. |
groupBy | groupBy(comptime K: type, gpa: Allocator, key: Series(K)) Error!GroupBy(K) | host | Sorted keys, so the grouping does not depend on row order. Gaps get a group. |
takeRows | takeRows(comptime T: type, gpa: Allocator, column: Series(T), wanted: []const usize) Error!Series(T) | host | ONE gather: filter, sort, sample and a join's assembly are index lists. |
firstRows | firstRows(comptime T: type, gpa: Allocator, column: Series(T), n: usize) Error!Series(T) | host | pandas' head. Not called `head` - that shadows three locals here. |
JoinHow | JoinHow = enum | host | inner, left, right, outer - znum's four with its meanings. |
RowPair | RowPair = struct | host | One joined row. `no_row` beats ?usize: dense array, tight loop. |
no_row | no_row: usize = zm.maxInt(usize); | method | The sentinel: nothing on this side. |
joinRows | joinRows(comptime K: type, gpa: Allocator, left_key: Series(K), right_key: Series(K), how: JoinHow) Error![]RowPair | host | The PAIRING without the gather. A missing key never matches, not even another. |
Summary | Summary(comptime T: type) type | host | NAMED fields. znum returns [8]f64 and slot 3 being min is not a type. |
describe | describe(comptime T: type, gpa: Allocator, column: Series(T)) Error!?Summary(T) | host | The pandas eight, skipping gaps. Null when nothing is present. |
Aggregate | Aggregate = enum | host | Seven that fit one pass. No variance or median stub that always errors. |
needsFloat | needsFloat(self: Aggregate) bool | method | Only mean does. znum floats all ten. |
aggregate | aggregate(comptime T: type, comptime how: Aggregate, gpa: Allocator, grouped: anytype, values: Series(T), name: []const u8) Error!Series(T) | host | The TYPE survives: sum of i32 is i32. A mean of i32 is a COMPILE error. |
factorize | factorize(comptime T: type, out_codes: []usize, out_levels: []T, a: Tensor(T), order: CodeOrder) Error!usize | host | Two orders, caller picks. Both defensible; the wrong one is silent. |
CodeOrder | CodeOrder = enum | host | first_seen like pandas, or sorted. Required. |
valueCounts | valueCounts(comptime T: type, out_values: []T, out_counts: []usize, scratch: []T, a: Tensor(T)) Error!usize | host | Distinct values and frequencies, most common first. Returns the count. |
logSumExp | logSumExp(comptime T: type, a: Tensor(T)) Error!T | host | Subtracts the max first, so exp cannot overflow. Direct: inf. This: finite. |
logSumExpAxis | logSumExpAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void | log_sum_exp_rows | Per row, so one saturating row does not affect its neighbours. |
stdDevAxis | stdDevAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize, ddof: Ddof) Error!void | host | The square root of varianceAxis, without a second pass. |
determinant | determinant(comptime T: type, a: Tensor(T), pivots: []usize) Error!T | host | Pivot sign times the diagonal. OVERWRITES `a`, as lu does. |
polyakUpdate | polyakUpdate(comptime T: type, target: Tensor(T), online: Tensor(T), follow: T) Error!void | host | Affine, not lerp: `follow = 1` is an exact copy. Not called `tau`. |
cartpoleStep | cartpoleStep(comptime T: type, state: CartpoleState(T), push: Push) CartpoleStep(T) | host | PURE: state in, next state out. znum mutates and warns you in prose. |
cartpoleStepBatch | cartpoleStepBatch(comptime T: type, next: Tensor(T), outcome: Tensor(T), state: Tensor(T), pushes: []const Push) Error!void | host | A LOOP over the scalar step - one definition, and the GPU twin is free. |
cartpoleReset | cartpoleReset(comptime T: type, rng: Rng, episode: u32) CartpoleState(T) | host | Separate from the step, which removes two flags. |
Push | Push = enum | cartpole_step | left or right. znum decodes `action > 0.5`. |
CartpoleState | CartpoleState(comptime T: type) type | host | Four NAMED fields, not four positions in a tensor. |
CartpoleStep | CartpoleStep(comptime T: type) type | host | state, reward, failed - not rd[0] and rd[1]. |
categorical | categorical(comptime T: type, logits: Tensor(T), rng: Rng, index: u32) Error!usize | host | One draw per sample. The max is subtracted, so logits of 800 work. |
boundedIndex | boundedIndex(rng: Rng, bound: usize, index: u32) Error!usize | host | Uniform without the modulo bias every tutorial ships. |
affine | affine(comptime T: type, out: Tensor(T), a: Tensor(T), gain: T, offset: T) Error!void | affine_f | a * gain + offset. One function where znum has three. |
divFloor | divFloor(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | -7/2 is -4. `div` stays float-only so integers must choose. |
divTrunc | divTrunc(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | -7/2 is -3. Both are correct; the caller says which. |
divCeil | divCeil(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | Rounds up. What you want when sizing a buffer. |
divExact | divExact(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | Refuses to round. @divExact is UNDEFINED when inexact; this errors. |
mod | mod(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | Sign follows the DIVISOR: -7 mod 3 is 2. Use this to wrap an index. |
rem | rem(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | Sign follows the DIVIDEND: -7 rem 3 is -1. |
logSoftmaxRows | logSoftmaxRows(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | Never forms the softmax: a value 90 below the max gives -90, not -inf. |
compareScalar | compareScalar(comptime T: type, out: Tensor(T), a: Tensor(T), value: T, how: Comparison) Error!void | host | One function where znum has six: eq/ne/lt/le/gt/geScalar. |
Comparison | Comparison = enum | host | Named, so the call site reads as a sentence. |
logicalAnd | logicalAnd(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | NOT bitwiseAnd: 2 & 1 is 0, logicalAnd(2, 1) is 1. |
logicalOr | logicalOr(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | One where either is nonzero. |
logicalXor | logicalXor(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | One where exactly one is nonzero. |
resolveInferredShape | resolveInferredShape(count: usize, dims: []const isize, out: []usize) Error!usize | host | The -1 dimension. Two of them, any other negative, or a non-divisor: error. |
swapaxes | swapaxes(comptime T: type, a: Tensor(T), first: usize, second: usize) Error!Tensor(T) | host | A VIEW - the strides trade places, no data moves. |
stack | stack(comptime T: type, out: Tensor(T), parts: []const Tensor(T), axis: usize) Error!void | host | A NEW axis, so rank grows. concat grows an EXISTING axis - the usual mixup. |
shiftRows | shiftRows(comptime T: type, gpa: Allocator, column: Series(T), by: isize) Error!Series(T) | host | The vacated rows are MISSING, not zero - zero would read as no change. |
rollingMean | rollingMean(comptime T: type, gpa: Allocator, column: Series(T), window: usize) Error!Series(T) | host | A partial window is MISSING. Averaging what is there looks artificially smooth. |
rfftfreq | rfftfreq(comptime T: type, out: Tensor(T), n: usize, spacing: T) Error!void | host | What frequency each bin IS. Without it a spectrum is unlabelled numbers. |
trapz | trapz(comptime T: type, y: Tensor(T), spacing: T) Error!T | host | ENDPOINTS count half - they belong to one interval, interiors to two. |
trapzCoords | trapzCoords(comptime T: type, y: Tensor(T), x: Tensor(T)) Error!T | host | Uneven spacing. A non-increasing x is refused, not silently negative area. |
PerClass | PerClass(comptime T: type) type | host | precision, recall, f1 AND support - support tells a measured 0 from a blank. |
perClassMetrics | perClassMetrics(comptime T: type, out: []PerClass(T), matrix: Tensor(T)) Error!void | host | Precision is a COLUMN sum, recall a ROW sum. F1 is the harmonic mean. |
histogramEdges | histogramEdges(comptime T: type, out: Tensor(T), low: T, high: T, count: usize) Error!void | host | n bins have n+1 edges. The last is set exactly, not computed. |
astype | astype(comptime From: type, comptime To: type, out: Tensor(To), a: Tensor(From)) Error!void | host | Truncates toward zero. Out of range is an ERROR: Zig @intFromFloat is UB. |
zeros | zeros(comptime T: type, gpa: Allocator, shape: []const usize) Error!Tensor(T) | host | alloc leaves GARBAGE - the common case gets a name. |
ones | ones(comptime T: type, gpa: Allocator, shape: []const usize) Error!Tensor(T) | host | As zeros, filled with one. |
full | full(comptime T: type, gpa: Allocator, shape: []const usize, value: T) Error!Tensor(T) | host | As zeros, filled with a value. |
fullLike | fullLike(comptime T: type, gpa: Allocator, model: Tensor(T), value: T) Error!Tensor(T) | host | Shaped like another. Writing a shape twice is how two drift apart. |
zerosLike | zerosLike(comptime T: type, gpa: Allocator, model: Tensor(T)) Error!Tensor(T) | host | Shaped like another, zero. |
onesLike | onesLike(comptime T: type, gpa: Allocator, model: Tensor(T)) Error!Tensor(T) | host | Shaped like another, one. |
eye | eye(comptime T: type, gpa: Allocator, rows: usize, columns: usize) Error!Tensor(T) | host | Identity, RECTANGULAR allowed - numpy does and it is useful. |
EndPoint | EndPoint = enum | host | include or exclude. Required: the wrong one is off by one. |
arange | arange(comptime T: type, gpa: Allocator, start: T, stop: T, step: T) Error!Tensor(T) | host | Stop EXCLUDED, as numpy, Python range and Zig 0..n all agree. |
linspace | linspace(comptime T: type, gpa: Allocator, start: T, stop: T, count: usize, endpoint: EndPoint) Error!Tensor(T) | host | Endpoint NAMED, and the last value is set exactly rather than computed. |
item | item(comptime T: type, a: Tensor(T)) Error!T | host | The one value of a one-element tensor. Refuses anything else. |
Extreme | Extreme = enum | host | largest or smallest, comptime - one branch, not a test. |
argExtreme | argExtreme(comptime T: type, comptime which: Extreme, a: Tensor(T)) Error!usize | host | FLAT index over the whole tensor. Ties go to the first, as numpy. |
all | all(comptime T: type, a: Tensor(T)) bool | all_nonzero | Empty is TRUE - forced by all(a ++ b) == all(a) and all(b). |
any | any(comptime T: type, a: Tensor(T)) bool | any_nonzero | Empty is FALSE, by the same argument on `or`. |
countNonzero | countNonzero(comptime T: type, a: Tensor(T)) usize | count_nonzero | A usize: a count is not in the tensor units. |
nanSum | nanSum(comptime T: type, a: Tensor(T)) T | host | A separate NAME, not a flag - sumAll propagating a NaN is useful. |
nanMean | nanMean(comptime T: type, a: Tensor(T)) T | host | NaN when every element is one. Series returns null. |
bitwiseAnd | bitwiseAnd(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | Bit by bit. Integer-only, checked at the front door. |
bitwiseOr | bitwiseOr(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | Bit by bit. |
bitwiseXor | bitwiseXor(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | Bit by bit. |
bitwiseNot | bitwiseNot(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | Flips every bit: ~5 is -6. NOT logicalNot, which asks if it was zero. |
logicalNot | logicalNot(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | One where the input was zero. znum has only this one. |
shiftLeft | shiftLeft(comptime T: type, out: Tensor(T), a: Tensor(T), places: usize) Error!void | host | An over-shift is a DomainError, not Zig's undefined behaviour. |
shiftRight | shiftRight(comptime T: type, out: Tensor(T), a: Tensor(T), places: usize) Error!void | host | Arithmetic when signed, so -8 >> 1 is -4. |
meshGrid | meshGrid(comptime T: type, out_x: Tensor(T), out_y: Tensor(T), x: Tensor(T), y: Tensor(T)) Error!void | host | Two coordinate grids from two axes. Outputs are NAMED, not a [2]Tensor. |
repeatEach | repeatEach(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize, count: usize) Error!void | repeat_each | [1,2,3] -> [1,1,2,2,3,3]. `tile` repeats the WHOLE tensor instead. |
moveAxis | moveAxis(comptime T: type, a: Tensor(T), axis: usize, destination: usize) Error!Tensor(T) | host | The other axes keep their order. |
cumprod | cumprod(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | Running product along the last axis. |
cummin | cummin(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | Running minimum. |
cummax | cummax(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | Running maximum. |
diff | diff(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | diff_forward | First position is NaN; the LENGTH IS KEPT so it still lines up. |
pctChange | pctChange(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | As diff, and NaN where the previous value is zero. |
argmaxAxis | argmaxAxis(comptime T: type, comptime Index: type, out: Tensor(Index), a: Tensor(T), axis: usize) Error!void | host | Index along an axis, as a float. NaN never wins. |
argminAxis | argminAxis(comptime T: type, comptime Index: type, out: Tensor(Index), a: Tensor(T), axis: usize) Error!void | host | Ties go to the first. |
layerNormRowsBackward | layerNormRowsBackward(comptime T: type, dx: Tensor(T), a: Tensor(T), dy: Tensor(T), epsilon: T) Error!void | host | Every element of a row appears in every other element's gradient. |
argminAll | argminAll(comptime T: type, a: Tensor(T)) Error!usize | host | Ties to the first. NaN never wins. |
flip | flip(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void | host | Reverses one axis. |
roll | roll(comptime T: type, out: Tensor(T), a: Tensor(T), shift: isize, axis: usize) Error!void | host | out[i] = a[(i - shift) mod n]. Positive moves forward. |
nonzero | nonzero(comptime T: type, out: []usize, a: Tensor(T)) Error!usize | host | Flat positions. Returns the count. |
booleanMask | booleanMask(comptime T: type, out: []T, a: Tensor(T), mask: Tensor(T)) Error!usize | host | The values at those positions. |
bincount | bincount(comptime T: type, counts: []usize, a: Tensor(T)) Error!void | host | Out of range is an error, not a silent drop. |
sort | sort(comptime T: type, out: []T, a: Tensor(T)) Error!void | host | Ascending, into a caller-supplied slice. |
unique | unique(comptime T: type, out: []T, a: Tensor(T)) Error!usize | host | Ascending, not first-seen. |
argsort | argsort(comptime T: type, out: []usize, a: Tensor(T)) Error!void | host | Rank 1, ascending, stable. Indices into a []usize. |
sumAxis | sumAxis(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void | sum_axis0 | The output rank is one less. |
matmul | matmul(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | matmul, matmul_tiled, matmul_bt | Rank 2. Compensated inner sum. |
relu | relu(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | reluz, relu | |
sigmoid | sigmoid(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | sigmoid | |
tanh | tanh(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | tanhf, tanhf | |
gelu | gelu(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | gelu | tanh approximation. |
materialise | materialise(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | slice_columns | Densifies any view. |
take | take(comptime T: type, out: Tensor(T), a: Tensor(T), indices: []const usize, axis: usize) Error!void | host | Gathers named positions. |
concat | concat(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T), axis: usize) Error!void | concat_columns | Differentiable. The backward is the slice, read backwards. |
tile | tile(comptime T: type, out: Tensor(T), a: Tensor(T), axis: usize) Error!void | host | Whole multiples only. |
minimum | minimum(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | minimum | min is a reserved word. |
maximum | maximum(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | maximum | |
sqrt | sqrt(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | sqrtf | NaN below zero. |
log | log(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | logf | NaN below zero, −inf at zero. |
floor | floor(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | floorf | |
ceil | ceil(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | ceilf | |
sign | sign(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | signf, signz | Zero maps to zero. |
square | square(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | square | Exact, unlike pow(a,2). |
reciprocal | reciprocal(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | reciprocal | Infinite at zero on both backends. |
trunc | trunc(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | truncf | Toward zero. |
round | round(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | roundf | Half away from zero. |
sinRad | sinRad(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | sinf | |
cosRad | cosRad(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | cosf | |
atan2Rad | atan2Rad(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | atan2f | Argument order is (y, x). |
hypot | hypot(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | hypotf | |
greater | greater(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | greater, where_pick | 1 or 0 in T, not a bool tensor. |
less | less(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | less | |
equal | equal(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | equal | Exact equality, no built-in epsilon. |
clamp | clamp(comptime T: type, out: Tensor(T), a: Tensor(T), lo: T, hi: T) Error!void | clampf | |
lerp | lerp(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T), t: T) Error!void | lerpf | Exact at both endpoints. |
sinTurns | sinTurns(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | sin_turns | Argument in TURNS. Every quarter turn exact; whole turns exactly zero. |
cosTurns | cosTurns(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | cos_turns | As sinTurns, a quarter turn ahead. |
tanTurns | tanTurns(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | INFINITE at 0.25 and 0.75: in turns the pole is reachable. |
tanRad | tanRad(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | |
asinRad | asinRad(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | CLAMPS outside [-1,1] — zm's choice. See its note. |
acosRad | acosRad(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | Clamps, like asin. |
atanRad | atanRad(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | |
sinh | sinh(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | sinhf | By identity, so CPU and GPU compute the same expression. |
cosh | cosh(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | coshf | |
asinh | asinh(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | asinhf | No znum counterpart. |
acosh | acosh(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | NaN below 1. No znum counterpart. |
atanh | atanh(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | atanhf | NaN outside (-1,1). No znum counterpart. |
rsqrt | rsqrt(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | rsqrtf | NaN below zero. |
pow | pow(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | exp(b·log a). NaN for a negative base. |
exp2 | exp2(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | exp2f | 2ⁿ. Named by the correspondence gate on its first run. |
log2 | log2(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | log2f | |
log10 | log10(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | log10f | |
expm1 | expm1(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | expm1, expm1z, expm1 | exp(x) − 1, without extra precision near zero. |
log1p | log1p(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | log1p, log1pz, log1p | log(1 + x), likewise. |
cbrt | cbrt(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | cbrtf | Defined for negatives: sign(x)·|x|^(1/3). |
notEqual | notEqual(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | not_equal | |
greaterEqual | greaterEqual(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | greater_equal | |
lessEqual | lessEqual(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | less_equal | |
scale | scale(comptime T: type, out: Tensor(T), a: Tensor(T), factor: T) Error!void | scale | Runtime factor, so not a zip. |
exp | exp(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | expf | |
abs | abs(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | absf | |
neg | neg(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | neg | |
reluGrad | reluGrad(comptime T: type, out: Tensor(T), x: Tensor(T), grad: Tensor(T)) Error!void | relu_grad | Takes the forward input. |
sigmoidGrad | sigmoidGrad(comptime T: type, out: Tensor(T), y: Tensor(T), grad: Tensor(T)) Error!void | sigmoid_grad | Takes the forward output. |
tanhGrad | tanhGrad(comptime T: type, out: Tensor(T), y: Tensor(T), grad: Tensor(T)) Error!void | tanh_grad | Takes the forward output. |
sgdStep | sgdStep(comptime T: type, out: Tensor(T), weight: Tensor(T), grad: Tensor(T), rate: T) Error!void | sgd_step | Pass weight as out for in place. |
softplus | softplus(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | softplus | Its OWN node: log(1+exp(x)) overflows at x=89 in f32. Backward is sigmoid. |
silu | silu(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | silu | x·sigmoid(x), also called swish. |
leakyRelu | leakyRelu(comptime T: type, out: Tensor(T), a: Tensor(T), slope: T) Error!void | leaky_relu | The slope has no default. |
elu | elu(comptime T: type, out: Tensor(T), a: Tensor(T), alpha: T) Error!void | elu | A real node: a branch is not a composition. Backward reads the INPUT. |
dotAll | dotAll(comptime T: type, a: Tensor(T), b: Tensor(T)) Error!T | host | Sizes must match, shapes need not. Compensated. |
trace | trace(comptime T: type, a: Tensor(T)) Error!T | host | Square, rank 2. |
norm | norm(comptime T: type, a: Tensor(T)) Error!T | host | Frobenius. Scaled, so it does not overflow. |
lu | lu(comptime T: type, a: Tensor(T), pivots: []usize) Error!T | host | In place. Partial pivoting. Returns the permutation sign. |
solve | solve(comptime T: type, a: Tensor(T), b: Tensor(T), pivots: []usize) Error!void | host | Overwrites a with its factors and b with x. Several right-hand sides at once. |
determinantLu | determinantLu(comptime T: type, a: Tensor(T), pivots: []usize) Error!T | host | Overwrites a. Returns 0 for a singular matrix. |
qr | qr(comptime T: type, q: Tensor(T), r: Tensor(T), a: Tensor(T)) Error!void | host | Householder. Orthogonality holds at 4e-16 where Gram-Schmidt reaches 1.8e-4. |
svd | svd(comptime T: type, u: Tensor(T), values: []T, v: Tensor(T), a: Tensor(T)) Error!void | host | One-sided Jacobi. Never forms aᵀa, so small singular values survive. |
pinv | pinv(comptime T: type, out: Tensor(T), u: Tensor(T), values: []T, v: Tensor(T), a: Tensor(T), rcond: T) Error!void | host | Via svd. Rank cutoff is rcond · sigma_max, so it is scale-free. |
inverse | inverse(comptime T: type, out: Tensor(T), a: Tensor(T), pivots: []usize) Error!void | host | By solving against the identity. Prefer solve. |
matmulNT | matmulNT(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | a · bᵀ, without materialising the transpose. |
matmulTN | matmulTN(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | aᵀ · b. |
bmm | bmm(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | Batched: (batch,m,k) @ (batch,k,n). |
tensordot | tensordot(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T), axis_a: usize, axis_b: usize) Error!void | host | The general product. matmul and dot are special cases. |
Named | Named(comptime T: type, comptime names: []const []const u8) type | host | Axis names in the TYPE. axis("width") is comptime; a typo is a compile error. |
axis_names | axis_names: []const []const u8 = names; | method | The axis names of a Named, in order. |
of | of(tensor: Tensor(T)) Error!Self | method | The tag for a float type. |
axis | axis(comptime name: []const u8) usize | method | Comptime index of a named axis. |
extent | extent(self: Self, comptime name: []const u8) usize | method | The length of a named axis. |
sumAxis | sumAxis(self: Self, comptime name: []const u8, out: Named(T, dropName(names, name))) Error!void | method | The output rank is one less. |
meanAxis | meanAxis(self: Self, comptime name: []const u8, out: Named(T, dropName(names, name))) Error!void | method | |
maxAxis | maxAxis(self: Self, comptime name: []const u8, out: Named(T, dropName(names, name))) Error!void | method | NaN never wins. |
window | window(self: Self, comptime name: []const u8, start: usize, len: usize) Error!Self | method | A mini-batch: same type back, so nothing downstream changes. |
transposed | transposed(self: Self, comptime first: []const u8, comptime second: []const u8) Error!Named(T, swapNames(names, first, second)) | method | Swaps two axes AND their names. |
namedMatmul | namedMatmul(comptime T: type, comptime a_names: []const []const u8, comptime b_names: []const []const u8, out: Named(T, &.{ a_names[0], b_names[1] }), a: Named(T, a_names), b: Named(T, b_names)) Error!void | host | The contracted axes must be the SAME axis, not just the same length. |
einsum | einsum(comptime T: type, comptime spec: []const u8, out: Tensor(T), operands: anytype) Error!void | host | Comptime spec, so a malformed one is a COMPILE error naming the problem. |
outer | outer(comptime T: type, out: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | out[i][j] = a[i] * b[j]. |
diagonal | diagonal(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | Rectangular too: min(rows, cols) entries. |
triangle | triangle(comptime T: type, out: Tensor(T), a: Tensor(T), side: Triangle, keep_diagonal: bool) Error!void | host | One triangle kept, the other zeroed. |
conditionNumber | conditionNumber(comptime T: type, u: Tensor(T), values: []T, v: Tensor(T), a: Tensor(T)) Error!T | host | sigma_max/sigma_min. Infinite when singular. |
eigvalsSymmetric | eigvalsSymmetric(comptime T: type, values: []T, scratch: Tensor(T), a: Tensor(T)) Error!void | host | Ascending, without keeping the vectors. |
Triangle | Triangle = enum | host | upper or lower. |
solveTriangular | solveTriangular(comptime T: type, x: Tensor(T), t: Tensor(T), b: Tensor(T), side: Triangle) Error!void | host | By substitution. DomainError on a zero diagonal. |
lstsq | lstsq(comptime T: type, x: Tensor(T), q: Tensor(T), r: Tensor(T), a: Tensor(T), b: Tensor(T)) Error!void | host | By QR. The normal equations square the condition number; this does not. |
eigh | eigh(comptime T: type, values: []T, vectors: Tensor(T), a: Tensor(T)) Error!void | host | Symmetric, Jacobi. Convergence relative to the norm, so scale-free. |
cholesky | cholesky(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | Lower-triangular L with L·Lᵀ = a. DomainError when not positive-definite. |
softmaxRows | softmaxRows(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | softmax_rows | Rank 2. The row maximum is subtracted. |
variance | variance(comptime T: type, a: Tensor(T), ddof: Ddof) Error!T | host | Two passes. ddof selects the divisor. |
stdDev | stdDev(comptime T: type, a: Tensor(T), ddof: Ddof) Error!T | host | |
rootMeanSquare | rootMeanSquare(comptime T: type, a: Tensor(T)) Error!T | host | |
meanAbsDev | meanAbsDev(comptime T: type, a: Tensor(T)) Error!T | host | Mean |x − mean|. |
stdErr | stdErr(comptime T: type, a: Tensor(T)) Error!T | host | Sample deviation over √n. |
zscore | zscore(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | Sample deviation. |
covariance | covariance(comptime T: type, a: Tensor(T), b: Tensor(T), ddof: Ddof) Error!T | host | |
correlation | correlation(comptime T: type, a: Tensor(T), b: Tensor(T)) Error!T | host | Not clamped to [−1, 1]. |
minMaxScale | minMaxScale(comptime T: type, out: Tensor(T), a: Tensor(T)) Error!void | host | DomainError on a constant tensor. |
layerNormRows | layerNormRows(comptime T: type, out: Tensor(T), a: Tensor(T), epsilon: T) Error!void | layernorm_rows | Rank 2, population variance, epsilon is a parameter. |
maeLoss | maeLoss(comptime T: type, a: Tensor(T), b: Tensor(T)) Error!T | mae_loss | An outlier contributes its distance, not its square. |
huberLoss | huberLoss(comptime T: type, a: Tensor(T), b: Tensor(T), delta: T) Error!T | huber_loss | Squared below delta, linear above. Slopes meet at the join. |
binaryCrossEntropyFromLogits | binaryCrossEntropyFromLogits(comptime T: type, logits: Tensor(T), targets: Tensor(T)) Error!T | bce_loss | Takes LOGITS. Finite at ±800, where the sigmoid form is not. |
klDivergenceRows | klDivergenceRows(comptime T: type, p: Tensor(T), q: Tensor(T)) Error!T | host | sum p·log(p/q). Not symmetric. |
argmaxAll | argmaxAll(comptime T: type, a: Tensor(T)) Error!usize | host | Flat index. Ties to the first; NaN never wins. |
argmaxRows | argmaxRows(comptime T: type, out: []usize, a: Tensor(T)) Error!void | argmax_axis0 | One column index per row, into a []usize. |
oneHotRows | oneHotRows(comptime T: type, out: Tensor(T), indices: []const usize) Error!void | host | Zeroes out first. |
crossEntropyRows | crossEntropyRows(comptime T: type, logits: Tensor(T), targets: []const usize) Error!T | host | Takes LOGITS. logsumexp form, finite for every finite input. |
crossEntropyRowsGrad | crossEntropyRowsGrad(comptime T: type, out: Tensor(T), logits: Tensor(T), targets: []const usize) Error!void | host | softmax − onehot, over the row count. |
confusionMatrix | confusionMatrix(comptime T: type, out: Tensor(T), predicted: []const usize, truth: []const usize) Error!void | host | True DOWN, predicted ACROSS. |
ClassScore | ClassScore = struct | host | Carries flags saying which zeros are undefined. |
classScore | classScore(comptime T: type, confusion: Tensor(T), k: usize) Error!ClassScore | host | Precision, recall, F1 for one class. |
macroF1 | macroF1(comptime T: type, confusion: Tensor(T)) Error!f64 | host | Unweighted, so a rare class counts as much as a common one. |
r2Score | r2Score(comptime T: type, prediction: Tensor(T), truth: Tensor(T)) Error!T | host | CAN BE NEGATIVE. Worse than the mean is information. |
cosineSimilarity | cosineSimilarity(comptime T: type, a: Tensor(T), b: Tensor(T)) Error!T | host | DomainError on a zero vector, not 0. |
hingeLoss | hingeLoss(comptime T: type, scores: Tensor(T), labels: Tensor(T)) Error!T | host | Margin 1: correct but close still costs. |
accuracy | accuracy(comptime T: type, logits: Tensor(T), targets: []const usize) Error!T | host | Fraction of rows whose largest logit is at the target. |
gradient | gradient(comptime T: type, out: []T, values: []const T, spacing: T) Error!void | host | Central inside, one-sided at the ends. |
trapezoid | trapezoid(comptime T: type, values: []const T, spacing: T) Error!T | host | Exact for any straight line. |
interp | interp(comptime T: type, out: []T, at: []const T, xs: []const T, ys: []const T) Error!void | host | Clamps outside the range rather than extrapolating. |
ComplexNumber | ComplexNumber(comptime T: type) type | host | extern, so an array is interleaved re/im. |
add | add(a: Self, b: Self) Self | method | |
sub | sub(a: Self, b: Self) Self | method | |
mul | mul(a: Self, b: Self) Self | method | Elementwise, not a matrix product. |
magnitude | magnitude(a: Self) T | method | |
dft | dft(comptime T: type, out: []ComplexNumber(T), input: []const ComplexNumber(T), exponent_sign: T) Error!void | host | O(n²), any length. The reference `fft` is checked against. |
fft | fft(comptime T: type, data: []ComplexNumber(T), exponent_sign: T) Error!void | host | Radix-2, in place, power-of-two. 630× faster than dft at 4096. |
ifft | ifft(comptime T: type, data: []ComplexNumber(T)) Error!void | host | The whole 1/n lives here. |
fftFreq | fftFreq(comptime T: type, out: []T, spacing: T) Error!void | host | Second half is negative frequencies. |
rfft | rfft(comptime T: type, out: []ComplexNumber(T), input: []const T, scratch: []ComplexNumber(T)) Error!void | host | Real input, n/2 + 1 bins. Half the transform, not a full one halved. |
ConvMode | ConvMode = enum | host | full, same, valid. Each is a window onto full. |
outputLen | outputLen(self: ConvMode, a: usize, v: usize) usize | method | Outputs a mode produces. `length` is reserved. |
convolve | convolve(comptime T: type, out: []T, a: []const T, v: []const T, extent: ConvMode) Error!void | host | Flips the kernel. |
correlate | correlate(comptime T: type, out: []T, a: []const T, v: []const T, extent: ConvMode) Error!void | host | Does not flip the kernel. |
magic | magic: []const u8 = "ZNUM\x00\x00"; | host | Six bytes. A wrong file is refused, not misread. |
format_version | format_version: u8 = 1; | host | Bumped when a layout change would be misread, not rejected. |
DType | DType = enum(u8) | host | One byte. An f32 file read as f64 is refused. |
byteWidth | byteWidth(self: DType) usize | method | Derived from the tag: 1 << (tag & 0xf). |
of | of(comptime T: type) DType | method | The tag for a float type. |
NamedTensor | NamedTensor = struct | host | A name, a shape, and erased bytes. |
saveNpy | saveNpy(comptime T: type, gpa: Allocator, out: *std.ArrayList(u8), a: Tensor(T)) Error!void | host | numpy's format, so other tools can read it. |
NpyHeader | NpyHeader = struct | host | What a .npy header said. |
readNpyHeader | readNpyHeader(bytes: []const u8) Error!NpyHeader | host | Shape, width and where the data starts. |
loadNpy | loadNpy(comptime T: type, out: Tensor(T), bytes: []const u8) Error!void | host | Refuses big-endian and Fortran order rather than mis-reading. |
Categorical | Categorical(comptime T: type) type | host | A TYPE so sample, logProb and entropy share one normalisation. |
sample | sample(self: Self, rng: Rng, row: usize, index: u32) Error!usize | method | Uniform with replacement. Addressable by (rng, index). |
logProb | logProb(self: Self, gpa: Allocator, row: usize, action: usize) Error!T | method | Via logSoftmax, not log(softmax) - a confident policy underflows the naive form. |
entropy | entropy(self: Self, gpa: Allocator, row: usize) Error!T | method | An exploration bonus. A Gaussian's does not depend on its mean. |
DiagGaussian | DiagGaussian(comptime T: type) type | host | log_std, not std: exp makes it positive free, and a clamp would zero its gradient. |
sample | sample(self: Self, rng: Rng, row: usize, out: []T) Error!void | method | Uniform with replacement. Addressable by (rng, index). |
logProb | logProb(self: Self, row: usize, action: []const T) Error!T | method | Via logSoftmax, not log(softmax) - a confident policy underflows the naive form. |
entropy | entropy(self: Self) T | method | An exploration bonus. A Gaussian's does not depend on its mean. |
SquashedGaussian | SquashedGaussian(comptime T: type) type | host | tanh into (-1,1). The correction via softplus, or it is -inf when saturated. |
log_std_min | log_std_min: T = -20; | method | Clamped in BOTH sample and logProb, from one constant. |
log_std_max | log_std_max: T = 2; | method | As log_std_min. SAC's usual values. |
sample | sample(self: Self, rng: Rng, row: usize, out_pre_squash: []T, out_action: []T) Error!void | method | Uniform with replacement. Addressable by (rng, index). |
logProb | logProb(self: Self, row: usize, pre_squash: []const T) Error!T | method | Via logSoftmax, not log(softmax) - a confident policy underflows the naive form. |
CriticAggregate | CriticAggregate = enum | host | One knob covers TD3, SAC, DroQ, REDQ and CrossQ. A SLICE, not a pair. |
aggregateCritics | aggregateCritics(comptime T: type, estimates: []const T, how: CriticAggregate) Error!T | host | A minimum of unbiased estimates is biased DOWN, and that is deliberate. |
randomSubset | randomSubset(gpa: Allocator, count: usize, subset: usize, rng: Rng) Error![]usize | host | REDQ's M of N, re-drawn so the pessimism is not systematic. |
sacTarget | sacTarget(comptime T: type, transition: Transition, next_estimates: []const T, next_log_prob: T, gamma: T, alpha: T, how: CriticAggregate) Error!T | host | The entropy term is what makes it SAC and not TD3. |
Temperature | Temperature = struct | host | log_alpha is the parameter: a clamp on alpha zeroes its own gradient. |
alpha | alpha(self: Temperature) f64 | method | exp(log_alpha) - the coefficient itself. |
gradient | gradient(self: Temperature, mean_log_prob: f64) f64 | method | Central inside, one-sided at the ends. |
DqnKind | DqnKind = enum | host | double fixes the max-of-noisy-estimates upward bias. |
Transition | Transition = struct | host | One step. terminal means no future to bootstrap. |
dqnTarget | dqnTarget(comptime T: type, transition: Transition, q_next_target: []const T, q_next_online: []const T, gamma: T, kind: DqnKind) Error!T | host | Online SELECTS, target EVALUATES - independent errors, so no bias feedback. |
explainedVariance | explainedVariance(comptime T: type, predicted: []const T, actual: []const T) Error!T | host | A critic predicting the mean has a small loss and explains NOTHING. |
PpoConfig | PpoConfig = struct | host | normalize_advantage ON by default - without it one rate does not transfer. |
PpoStats | PpoStats = struct | host | clip_fraction and approx_kl are how you know it works. The loss alone is not. |
ppoObjective | ppoObjective(advantages: []f64, log_prob_old: []const f64, log_prob_new: []const f64, entropies: []const f64, config: PpoConfig) Error!PpoStats | host | KL via exp(r)-1-r, which is NON-NEGATIVE. The cheap estimator can go below 0. |
RolloutStep | RolloutStep = struct | host | terminal kills the BOOTSTRAP; episode_end kills the TRACE. Not one flag. |
gae | gae(steps: []const RolloutStep, gamma: f64, lambda: f64, out_advantages: []f64, out_returns: []f64) Error!void | host | lambda 0 is the one-step error, 1 is the full return. Runs backward. |
clipGradNorm | clipGradNorm(comptime T: type, graph: *const Graph(T), parameters: []const Var, max_norm: T) Error!T | host | ONE norm over all parameters - per-tensor clipping changes the DIRECTION. |
saveParameters | saveParameters(comptime T: type, gpa: Allocator, out: *std.ArrayList(u8), graph: *const Graph(T), parameters: []const Var) Error!void | host | A trained network to bytes. Named p0, p1 - so the ORDER must match. |
loadParameters | loadParameters(comptime T: type, gpa: Allocator, bytes: []const u8, graph: *const Graph(T), parameters: []const Var) Error!void | host | Back into a network of the same shape. loadTensors refuses a mismatch. |
saveTensors | saveTensors(comptime T: type, gpa: Allocator, out: *std.ArrayList(u8), names: []const []const u8, tensors: []const Tensor(T)) Error!void | host | Little-endian, explicit widths. Contiguous only. |
loadTensors | loadTensors(comptime T: type, bytes: []const u8, names: []const []const u8, into: []const Tensor(T)) Error!void | host | Into tensors the caller holds, matched by name. |
generalizedAdvantage | generalizedAdvantage(comptime T: type, rewards: []const T, values: []const T, next_values: []const T, terminal: []const T, episode_end: []const T, gamma: T, lambda: T, out_advantages: []T, out_returns: []T) Error!void | host | terminal stops the bootstrap; episode_end resets the running sum. |
discountedReturns | discountedReturns(comptime T: type, rewards: []const T, terminal: []const T, gamma: T, out: []T) Error!void | host | GAE with lambda 1 and zero values. |
quantile | quantile(comptime T: type, a: Tensor(T), q: T, scratch: []T) Error!T | host | Linear interpolation at q·(n−1), as numpy does. |
quantileSorted | quantileSorted(comptime T: type, sorted: []const T, q: T) Error!T | host | For several quantiles of one sample. |
median | median(comptime T: type, a: Tensor(T), scratch: []T) Error!T | host | quantile(0.5). |
skew | skew(comptime T: type, a: Tensor(T)) Error!T | host | Population form. Zero for anything symmetric. |
kurtosis | kurtosis(comptime T: type, a: Tensor(T)) Error!T | host | EXCESS: three subtracted, so normal reads as zero. |
geoMean | geoMean(comptime T: type, a: Tensor(T)) Error!T | host | exp(mean(log x)). The product form overflows. |
harmMean | harmMean(comptime T: type, a: Tensor(T)) Error!T | host | The right average for rates. |
histogram | histogram(comptime T: type, counts: []usize, a: Tensor(T), low: T, high: T) Error!void | host | Equal-width bins, top edge closed. |
normalizeAdvantages | normalizeAdvantages(comptime T: type, advantages: []T, epsilon: T) Error!void | host | In place. Epsilon, not DomainError, for a constant batch. |
ppoClipObjective | ppoClipObjective(comptime T: type, logp_new: []const T, logp_old: []const T, advantages: []const T, clip: T) Error!T | host | An objective to MAXIMISE. min(ratio·A, clip(ratio)·A). |
entropyRows | entropyRows(comptime T: type, out: []T, probs: Tensor(T)) Error!void | host | A zero probability contributes zero. |
ReplayBuffer | ReplayBuffer(comptime T: type) type | host | A ring: the oldest transition is overwritten. |
init | init(gpa: Allocator, capacity: usize, width: usize) Error!Self | method | Rng.init(seed) and Ctx.init(gpa, seed). |
push | push(self: *Self, transition: []const T) Error!void | method | |
sample | sample(self: *const Self, rng: Rng, index: u32, out: []T) Error!void | method | Uniform with replacement. Addressable by (rng, index). |
conv2d | conv2d(comptime T: type, out: Tensor(T), image: Tensor(T), kernel: Tensor(T), stride: [2]usize, pad: [2]usize) Error!void | conv2d_same | Cross-correlation with stride and zero padding, per axis. Rank 2. |
maxPool2d | maxPool2d(comptime T: type, out: Tensor(T), image: Tensor(T), size: usize) Error!void | max_pool2d | Non-overlapping square windows. Partial trailing window dropped. |
avgPool2d | avgPool2d(comptime T: type, out: Tensor(T), image: Tensor(T), size: usize) Error!void | avg_pool2d | |
initXavier | initXavier(comptime T: type, w: Tensor(T), rng: Rng, fan_in: usize, fan_out: usize) Error!void | host | Uniform ±√(6/(in+out)). For tanh and sigmoid. |
initHe | initHe(comptime T: type, w: Tensor(T), rng: Rng, fan_in: usize) Error!void | host | Normal, deviation √(2/in). The 2 is relu's discarded half. |
InitScheme | InitScheme = enum | host | xavier or he. Named once, used by every layer. |
Conv2d | Conv2d(comptime T: type) type | host | Learned kernel and scalar bias. Dense's template. |
Attached | Attached = struct | method | What attach returns: the output plus both parameter handles. |
init | init(gpa: Allocator, rng: Rng, size: [2]usize, stride: [2]usize, pad: [2]usize, scheme: InitScheme) Error!Self | method | Rng.init(seed) and Ctx.init(gpa, seed). |
attach | attach(self: Self, graph: *Graph(T), input: Var) Error!Attached | method | Records input @ weight + bias onto a graph. |
TransformerBlock | TransformerBlock(comptime T: type) type | host | PRE-norm: the residual path is clear, so no warmup is needed. |
Attached | Attached = struct | method | What attach returns: the output plus both parameter handles. |
FeedForward | FeedForward = enum | method | relu or swiglu. Chosen at init: the gate changes widen's SHAPE. |
init | init(gpa: Allocator, rng: Rng, model_dim: usize, hidden: usize, heads: usize, epsilon: T, feed_forward: FeedForward) Error!Self | method | Rng.init(seed) and Ctx.init(gpa, seed). |
attach | attach(self: Self, graph: *Graph(T), input: Var, look: Attention(T).Look) Error!Attached | method | Records input @ weight + bias onto a graph. |
Attention | Attention(comptime T: type) type | host | Whether it can see the future is an ARGUMENT. Scale from the shape. |
Look | Look = enum | method | everywhere or backward_only. Required at every call. |
Attached | Attached = struct | method | What attach returns: the output plus both parameter handles. |
init | init(gpa: Allocator, rng: Rng, model_dim: usize, width: usize, heads: usize) Error!Self | method | Rng.init(seed) and Ctx.init(gpa, seed). |
attach | attach(self: Self, graph: *Graph(T), input: Var, look: Look) Error!Attached | method | Records input @ weight + bias onto a graph. |
LstmCell | LstmCell(comptime T: type) type | host | The forget bias starts at ONE: 44x the memory after ten steps. |
Gate | Gate = struct | method | sigmoid or silu, chosen at comptime. |
Stepped | Stepped = struct | method | hidden AND cell - two states, and confusing them is the classic bug. |
init | init(gpa: Allocator, rng: Rng, input_size: usize, hidden_size: usize) Error!Self | method | Rng.init(seed) and Ctx.init(gpa, seed). |
step | step(self: Self, graph: *Graph(T), input: Var, hidden: Var, cell: Var) Error!Stepped | method | One timestep, built from tape ops so the gradient is free. |
BatchNorm | BatchNorm(comptime T: type) type | host | The mode is an ARGUMENT, not a field - it cannot be forgotten. |
Use | Use = enum | method | training or inference. Required at every attach. |
Attached | Attached = struct | method | What attach returns: the output plus both parameter handles. |
init | init(gpa: Allocator, features: usize, momentum: T, epsilon: T) Error!Self | method | Rng.init(seed) and Ctx.init(gpa, seed). |
attach | attach(self: Self, graph: *Graph(T), input: Var, use: Use) Error!Attached | method | Records input @ weight + bias onto a graph. |
observe | observe(self: Self, batch: Tensor(T)) Error!void | method | Moves the running statistics. Separate, because attach only reads. |
Chain | Chain(comptime T: type, comptime Layers: type) type | host | Layers in order, parameters COLLECTED - a missing one never trains. |
Parameters | Parameters = [parameterCount()]Var; | method | Fixed array, sized at comptime from the layer list. |
Run | Run = struct | method | What Chain.attach returns: the output and every parameter. |
attach | attach(self: Self, graph: *Graph(T), input: Var) Error!Run | method | Records input @ weight + bias onto a graph. |
Embedding | Embedding(comptime T: type) type | host | A table of learned rows, indexed by whole numbers. |
Attached | Attached = struct | method | What attach returns: the output plus both parameter handles. |
init | init(gpa: Allocator, rng: Rng, count: usize, width: usize) Error!Self | method | Rng.init(seed) and Ctx.init(gpa, seed). |
attach | attach(self: Self, graph: *Graph(T), rows: []const usize) Error!Attached | method | Records input @ weight + bias onto a graph. |
Dense | Dense(comptime T: type) type | host | Owns its weights. Holds no Var, so it can attach to any graph. |
init | init(gpa: Allocator, rng: Rng, fan_in: usize, fan_out: usize, scheme: InitScheme) Error!Self | method | Rng.init(seed) and Ctx.init(gpa, seed). |
Attached | Attached = struct | method | What attach returns: the output plus both parameter handles. |
attach | attach(self: Self, graph: *Graph(T), input: Var) Error!Attached | method | Records input @ weight + bias onto a graph. |
gatherRows | gatherRows(comptime T: type, out: Tensor(T), table: Tensor(T), rows: []const usize) Error!void | host | out[i] = table[rows[i]]. |
scatterAddRows | scatterAddRows(comptime T: type, table: Tensor(T), source: Tensor(T), rows: []const usize) Error!void | host | table[rows[i]] += source[i]. Repeats accumulate. |
Optimizer | Optimizer(comptime T: type) type | host | Owns per-parameter state. One type, three algorithms. |
Kind | Kind = union(enum) | method | sgd, momentum, or adam — chosen at one call site. |
init | init(gpa: Allocator, graph: *const Graph(T), params: []const Var, kind: Kind) Error!Self | method | Rng.init(seed) and Ctx.init(gpa, seed). |
step | step(self: *Self, graph: *const Graph(T)) Error!void | method | One timestep, built from tape ops so the gradient is free. |
Var | Var = u32; | host | A handle into a Graph. An index, not a pointer. |
Graph | Graph(comptime T: type) type | host | Reverse-mode autograd. Eager forward, recorded backward. |
init | init(gpa: Allocator) Self | method | Rng.init(seed) and Ctx.init(gpa, seed). |
parameter | parameter(self: *Self, value: Tensor(T)) Error!Var | method | A value the graph differentiates with respect to. |
constant | constant(self: *Self, value: Tensor(T)) Error!Var | method | A fixed value; accumulates no gradient. |
valueOf | valueOf(self: *const Self, v: Var) Tensor(T) | method | The tensor behind a Var. |
gradOf | gradOf(self: *const Self, v: Var) Error!Tensor(T) | method | DomainError for a constant. |
add | add(self: *Self, a: Var, b: Var) Error!Var | method | |
matmul | matmul(self: *Self, a: Var, b: Var) Error!Var | method | Rank 2. Compensated inner sum. |
sub | sub(self: *Self, a: Var, b: Var) Error!Var | method | |
mul | mul(self: *Self, a: Var, b: Var) Error!Var | method | Elementwise, not a matrix product. |
relu | relu(self: *Self, a: Var) Error!Var | method | |
sigmoid | sigmoid(self: *Self, a: Var) Error!Var | method | |
tanh | tanh(self: *Self, a: Var) Error!Var | method | |
exp | exp(self: *Self, a: Var) Error!Var | method | |
log | log(self: *Self, a: Var) Error!Var | method | NaN below zero, −inf at zero. |
softplus | softplus(self: *Self, a: Var) Error!Var | method | Its OWN node: log(1+exp(x)) overflows at x=89 in f32. Backward is sigmoid. |
elu | elu(self: *Self, a: Var, alpha: T) Error!Var | method | A real node: a branch is not a composition. Backward reads the INPUT. |
mish | mish(self: *Self, x: Var) Error!Var | method | x * tanh(softplus(x)) - a composition once softplus is a node. |
silu | silu(self: *Self, x: Var) Error!Var | method | x·sigmoid(x), also called swish. |
gelu | gelu(self: *Self, x: Var) Error!Var | method | tanh approximation. |
glu | glu(self: *Self, x: Var) Error!Var | method | slice * sigmoid(slice). No backward: znum hand-derives one. |
swiglu | swiglu(self: *Self, x: Var) Error!Var | method | The same, gated by SiLU. Every recent transformer uses it. |
Gate | Gate = enum | method | sigmoid or silu, chosen at comptime. |
ppoClipLoss | ppoClipLoss(self: *Self, logp_new: Var, logp_old: []const T, advantages: []const T, clip: T) Error!Var | method | PPO as a LOSS. The clip acts in the backward: a clipped step gets ZERO. |
crossEntropy | crossEntropy(self: *Self, logits: Var, labels: []const usize) Error!Var | method | Labels are a []usize, not a Var: they have no gradient. |
softmaxRows | softmaxRows(self: *Self, a: Var) Error!Var | method | Rank 2. The row maximum is subtracted. |
layerNormRows | layerNormRows(self: *Self, a: Var, epsilon: T) Error!Var | method | Rank 2, population variance, epsilon is a parameter. |
conv2d | conv2d(self: *Self, input: Var, kernel: Var, stride: [2]usize, pad: [2]usize) Error!Var | method | Cross-correlation with stride and zero padding, per axis. Rank 2. |
slice | slice(self: *Self, a: Var, axis: usize, start: usize, len: usize) Error!Var | method | A window. Shares storage. |
concat | concat(self: *Self, a: Var, b: Var, axis: usize) Error!Var | method | Differentiable. The backward is the slice, read backwards. |
embedding | embedding(self: *Self, table: Var, rows: []const usize) Error!Var | method | On the graph. Gradient is a scatter-add. |
transpose | transpose(self: *Self, a: Var) Error!Var | method | Exchanges two axes. |
scale | scale(self: *Self, a: Var, factor: T) Error!Var | method | Runtime factor, so not a zip. |
dropout | dropout(self: *Self, a: Var, rate: T, rng: Rng) Error!Var | method | Inverted. Mask drawn once and stored; see resampleDropout. |
resampleDropout | resampleDropout(self: *Self, rng: Rng) void | method | Fresh masks for every dropout node. Once per step. |
mseLoss | mseLoss(self: *Self, a: Var, target: Var) Error!Var | method | Returns a scalar. Shapes must match exactly. |
backward | backward(self: *Self, loss: Var) Error!void | method | Seeds 1 into a single-element value and walks the tape back. |
recompute | recompute(self: *Self) Error!void | method | Replays the tape forward. Leaves keep whatever was written. |
checkGradient | checkGradient(self: *Self, loss: Var, param: Var, h: T) Error!T | method | Worst gap against a central finite difference. Restores the model. |
sgdMomentum | sgdMomentum(comptime T: type, out: Tensor(T), weight: Tensor(T), grad: Tensor(T), velocity: Tensor(T), rate: T, momentum: T) Error!void | sgd_momentum | velocity updated in place. momentum 0 equals sgdStep. |
adamStep | adamStep(comptime T: type, out: Tensor(T), weight: Tensor(T), grad: Tensor(T), moment: Tensor(T), velocity: Tensor(T), hyper: Adam, step: usize) Error!void | adam_step | moment and velocity updated in place. step is 1-based, for the bias correction. |
adamWStep | adamWStep(comptime T: type, out: Tensor(T), weight: Tensor(T), grad: Tensor(T), moment: Tensor(T), velocity: Tensor(T), hyper: Adam, decay: T, step: usize) Error!void | host | Decay applied to the WEIGHT, not the gradient. |
rmspropStep | rmspropStep(comptime T: type, out: Tensor(T), weight: Tensor(T), grad: Tensor(T), mean_square: Tensor(T), rate: T, rho: T, epsilon: T) Error!void | host | Adam without the first moment or bias correction. |
adagradStep | adagradStep(comptime T: type, out: Tensor(T), weight: Tensor(T), grad: Tensor(T), accumulator: Tensor(T), rate: T, epsilon: T) Error!void | host | Never forgets, so the rate only shrinks. |
clipByValue | clipByValue(comptime T: type, out: Tensor(T), a: Tensor(T), limit: T) Error!void | host | Elementwise, so it CHANGES direction. See clipByNorm. |
cosineLearningRate | cosineLearningRate(comptime T: type, step: usize, total: usize, base: T, lowest: T) T | host | base down to lowest, steepest in the middle. |
warmupLearningRate | warmupLearningRate(comptime T: type, step: usize, warmup: usize, base: T) T | host | Linear ramp. 0 at step 0, base at step warmup. |
stepLearningRate | stepLearningRate(comptime T: type, step: usize, every: usize, base: T, gamma: T) Error!T | host | A staircase: flat, then a drop. |
exponentialLearningRate | exponentialLearningRate(comptime T: type, step: usize, base: T, decay: T) T | host | base * exp(-decay * step). |
Adam | Adam = struct | host | Hyperparameters, defaulting to the paper's values. |
clipByNorm | clipByNorm(comptime T: type, out: Tensor(T), a: Tensor(T), limit: T) Error!T | host | Scales the whole tensor. Returns the norm before clipping. |
mseLoss | mseLoss(comptime T: type, prediction: Tensor(T), target: Tensor(T)) Error!T | mse_loss | Returns a scalar. Shapes must match exactly. |
Walk | Walk = struct | mesh_grid_y | Every position of a shape, in order. Replaces a five-line pattern written 48 times. |
over | over(shape: []const usize) Walk | method | Start a Walk over a shape. |
next | next(self: *Walk) ?[]const usize | method | The next position, or null. |
Walk2 | Walk2 = struct | host | Two shapes of equal size stepped together: the ranks may differ. |
Pair | Pair = struct | method | The two positions of one Walk2 step. |
over | over(first_shape: []const usize, second_shape: []const usize) Walk2 | method | Start a Walk over a shape. |
next | next(self: *Walk2) ?Pair | method | The next position, or null. |