Float32-in-ByteArray utilities.
All tensor data (params, images, gradients) stored as raw float32 bytes
in `ByteArray`. Zero conversion at the FFI boundary — IREE sees the same
bytes Lean wrote. Lean-side reads widen to `Float` (f64) only when needed
(loss printing, argmax, debugging).
Heavy-lift operations (He init, const fill, image loading) are @[extern]
to C for speed — avoids millions of Lean-level push calls.
Tile one flattened image base[off .. off+d0) (element offset off, d0 floats)
into m copies, each with independent N(0, σ²) exact Gaussian noise added
(Box-Muller). Returns m*d0 float32. NO clipping — the randomized-smoothing
certificate (Cohen–Rosenfeld–Kolter 2019) lives in raw input L2 space. With m=1,
off=0, d0=bs·pix it noises a whole training batch (Gaussian data augmentation).
Perturb one image base[off .. off+d0) by r·u for a uniformly-random unit vector u
(so ‖r·u‖₂ = r exactly). Returns the d0-vector x + r·u. For the Lipschitz-hypothesis
probe: shift the input by a known L2 amount and watch Φ⁻¹(P[f(x+η)=c]) respond.
Write three consecutive f32 values starting at float index idx, in place
when the array is unshared (it copies otherwise, so the result never depends
on a refcount). Used to patch the lr/bc₁/bc₂ slots of the Adam step
buffer without rebuilding it — see planning/archive/xla_pjrt_ladder.md §8.
n independent Bernoulli(keep) draws, survivors scaled 1/keep, as n float32 —
the hot loop of F32.dropoutMask, and nothing else. The guards (keep ≥ 1, n = 0,
n < 3) stay with the caller; entering here always means a real draw.
⚠ It is @[extern] for the reason this file's header gives — "avoids millions of
Lean-level push calls" — and that reason turned out to be quantitative rather than
stylistic: see dropoutMask below and ffi/f32_helpers.c.
dst[dstOff + i] += a · src[srcOff + i] for i < count, in place when dst is
unshared. The perturbation primitive of the adjoint gradcheck
(tests/TestR50GradCheck.lean): parameter tensors are packed in func-arg order, so a
direction supported on one BLOCK is a contiguous slice and needs no mask.
Σ_{i<count} a[aOff+i] · b[bOff+i], accumulated in f64.
⚠ The wide accumulator is the point. This computes the predicted directional derivative
⟨g, δ⟩ over as many as 4.7M same-sign terms; an f32 running sum would lose ~log₂ n bits to
the accumulation and report it as disagreement between the gradient and the finite
difference — i.e. as a failure of the thing under test. Returns NaN on an out-of-range
slice rather than clamping, so a caller that mis-computes an offset sees it.
Per-site, per-example stochastic-depth scales — bernoulli(keep_i)/keep_i for each of
keeps.size drop sites × bs examples, laid out site-major to match the render's
%dp<i>: tensor<Bxf32> inputs in signature order. Returns keeps.size * bs float32.
Pure Lean on purpose, in two ways.
Not in the graph. stablehlo.rng is disqualified: every numeric gate in this repo is a
bit-exactness or known-answer argument over a DETERMINISTIC graph — the tie harnesses' A-vs-A
floor, residency_gate.sh's bit-identity, the duplicated-batch DP identity, the cross-lowerer
IREE-vs-XLA agreement. A graph that draws its own randomness makes each of those either
impossible or contingent on seeding an XLA RNG identically across two lowerers and two vendors.
Not in C either. This is keeps.size * bs floats per step — 288 at EfficientNet's 9 sites and
batch 32, against a ~310 ms step — so the C round trip buys nothing measurable, and keeping the
draw in Lean keeps the one piece of genuine randomness in the training loop readable and seeded
where it can be audited. heInit is extern because it fills millions of values; this does not.
⚠ seed must be derived from the GLOBAL STEP, like augSeed, or no run is reproducible and
every gate that replays a step breaks. ⚠ 1/keep is folded in HERE rather than baked into the
graph — see VerifiedNet.dropKeeps. At keep = 1 the scale is exactly 1.0 for every example,
so a site with keep = 1 is the identity in IEEE, not merely close.
Equations
- One or more equations did not get rendered due to their size.
Instances For
▶ The classifier-dropout mask (recipe_gaps.md gap C) — n INDEPENDENT Bernoulli draws at
keep probability keep, each survivor scaled by 1/keep. Fills one %do: tensor<B×w×f32>
graph input, so n = B * w.
⚠⚠ n DRAWS, NOT B — this is dropScales' per-example loop replaced by a per-ELEMENT one,
and that single difference is the whole distinction between the two regularisers. The
reference draws bernoulli(key, keep, x.shape) for the classifier
(jax/Jax/Codegen.lean:1971) against (branch.shape[0],) + (1,)*(ndim-1) for stochastic depth
(:1037). A mask built by drawing B values and repeating each w times type-checks, fills
the same buffer, trains and descends — it is stochastic depth on the classifier. Nothing
downstream can tell: the shapes agree, the emitted graph is identical, and only the
DISTRIBUTION differs. Proofs.dropPath_scales_uniformly is the statement of what would be
wrong; on this side it is the loop bound, and there is no gate but reading it.
⚠ Seed it from a stream DISJOINT from dropScales'. The reference offsets by 999983
(fold_in(drop_key, 999983) against fold_in(drop_key, block_index)) exactly so a net running
both regularisers does not correlate them; the caller adds that offset. Sharing a seed here
would make the classifier mask a function of block 0's drop decisions, every step.
⚠ Same 1/keep folding as dropScales, for the same reason: the graph bakes no constant, so a
ones mask (keep ≥ 1, or eval) is the EXACT identity rather than a rescale.
Drawn on the HOST, not in the graph, for dropScales' reasons — stablehlo.rng would make
every bit-exactness gate in the repo contingent on seeding an XLA RNG identically across two
lowerers. That stays true and is not what changed here.
⛔⛔ WHAT CHANGED, AND WHAT THE OLD NOTE GOT WRONG. This used to push n boxed Floats into
an Array Float and tile them out through n/3 write3 calls, under the note "it is bigger
than dropScales (40,960 floats at B=32×1280, against 288) but still ~2 ULP of a ~310 ms
step". Both halves were wrong at the shape that matters. EfficientNet-B0's ImageNet job runs
global batch 256, so n = 256 × 1280 = 327,680 — 8× the shape that estimate was written
for — and measured against the compiled objects it cost 150.07 ms PER STEP:
| per step, measured 2026-08-30 | ms |
|---|---|
| this function, enet job shape (256×1280) | 150.07 |
| this function at 32×1280, the costed shape | 18.74 |
dropScales (9 sites × 256) | 1.83 |
F32.const, same 327,680 floats, @[extern] C | 0.073 |
It was 62% of B0's 281 ms step — more than twice its 71.6 ms graph — and B0 is the only
net that pays it (dropoutKeep is set on the two EfficientNet specs and nowhere else), which
is exactly why that net read as 2.63× its JAX reference when the others sit near parity.
▶ The loop now lives in dropoutFill/ffi/f32_helpers.c, byte-identical; the guards and the
seeding argument stay here, where they can be audited.
⭐ The transferable lesson: the estimate was not merely stale, it was taken at a shape no production job runs. A host-side cost is a function of the batch, so cost it at the batch.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Concatenate multiple ByteArrays. Fast (memcpy per chunk).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Extract the loss (last float32) from a train_step output.
Equations
- F32.extractLoss out lossIdx = F32.read out lossIdx.toUSize
Instances For
Drop the trailing loss float from train_step output.
Equations
- F32.dropLoss out nParams = out.extract 0 (nParams * 4)
Instances For
Decode the little-endian int32 label at record i of a packed label buffer.
⚠ Replaces lbl.get! (4 * i), which returned a UInt8 — byte 0 only, i.e. label % 256.
Every net this repo GATES is 10-class (Imagenette / CIFAR / MNIST), so byte 0 is the label
there and the truncation was invisible; on 1000-class ImageNet it silently discarded the high
byte, and since a correct prediction can then only match on classes 0..255, it capped every
reported top-1 at roughly a quarter of the truth. Measured off the val wire 2026-08-05: the
first batch carries labels 1..988 with 193 of 256 (75.4%) above 255.
Equations
Instances For
Argmax over n float32 values starting at element offset off.
⚠ Replaced argmax10, which took no n and scanned a literal 10 entries. Every net this
repo GATES is 10-class (Imagenette / CIFAR / MNIST), so the constant was right everywhere it
was ever checked and wrong on exactly the un-gated tier — the 1000-class ImageNet trainers,
where it confined every prediction to labels 0..9. Pass the net's own class count at the call
site; it is already the multiplier in the off expression, so the two cannot disagree.
Rank of label in a row of n logits at element offset off: the count of entries
STRICTLY GREATER than the label's own logit. The label is in the top-k iff rank < k.
Deliberately the same construction as the JAX reference's
jnp.sum(logits > true_logit, axis=1) < 5 rather than a sort or top_k — that side records
jax.lax.top_k's indices as broken on ROCm/gfx1100, and matching the formulation makes the
two paths' top-5 comparable by construction, ties included (strictly-greater means a tie
resolves in the label's favour on both sides).
Slice a batch of images: count images × pixelsPerImage floats. Zero-copy.
Equations
Instances For
sliceImages, zero-padding past the end of the dataset (total images).
The batch dimension is baked into the compiled .vmfb, so a final partial
eval batch must still be a full count images; the caller scores only the
first total - start rows of the result.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Slice a batch of labels: count records of bytesPerLabel bytes
each. Defaults to 4 (int32 LE) for classification. Per-pixel
segmentation masks pass bytesPerLabel := H * W (e.g. 224*224 = 50176
for Pets). Zero-copy.
Equations
Instances For
Load Oxford-IIIT Pets binary file. Returns (images f32 ByteArray, masks uint8 ByteArray, count). Images are 224×224×3, channel-first, normalized with ImageNet mean/std. Masks are 224×224 uint8 per-pixel class labels (0=fg, 1=bg, 2=boundary).
Load a BraTS (MSD Task01_BrainTumour) binary file at the given in-plane
size. Returns (images f32 ByteArray, masks uint8 ByteArray, count).
Images are imgSize×imgSize×4 (FLAIR / T1w / T1gd / T2w), channel-first.
Unlike the RGB datasets these carry no ImageNet normalization: the loader
inverts the uint8 quantization preprocess_brats.py applied, yielding the
per-volume, per-modality z-scored intensities the preprocessor computed over
brain voxels. Masks are imgSize×imgSize uint8 per-pixel class labels
(0=background, 1=edema, 2=non-enhancing tumour, 3=enhancing tumour).
YOLOv1 detection-bin loader (target+mask format; used by Pets). Returns (images_f32_normalized, yLabels_concat, count) where yLabels_concat carries 7200 bytes
per image: 30×7×7 float32 target (5880), then 7×7 float32 mask (196),
then numBoxes (4), then raw_boxes 56×20 (1120) — the Phase 3b format,
matching petsDetIO.labelBytesPerRecord. (This docstring said 6076,
the pre-Phase-3b target+mask size, long after the record grew the bbox
tail; a stale stride in the docs is what this whole bug class feeds on.)
The Lean dispatcher (runTraining) splits this into target + mask before
calling trainStepAdamF32Yolov1. See preprocess_pets_mosaic.py for the
on-disk format.
Dimension-parameterized detection-bin loader (same record format as
loadDetBin, but for an arbitrary square input imgSize and grid
gridH×gridW). Used for the higher-resolution VisDrone path (448 input /
14×14 grid); loadDetBin is the fixed 224/7×7 Pets path.
Anchor-format detection loader (brick #2). Returns (images_f32_normalized, target_only_concat, count) with numAnchors·15·gridH·gridW f32 target per
image — the anchor loss derives its per-anchor mask from the target's
objectness channels, so the on-disk mask/numBoxes/raw_boxes are skipped.
FPN multi-scale detection loader (brick #3). Returns (images_f32_normalized, target_only_concat, count) with ntot f32 target per image — the flat
[P3|P4|P5] block (ntot = Σ_s numAnchorsₛ·15·g_s²). Like the anchor loader,
the loss derives per-anchor masks from the target's objectness channels, so
the on-disk record is just image + flat target (no mask/boxes). Eval GT comes
from the single-box val.bin geometry, as in the anchor path.
Split an interleaved YOLOv1 batch slice into separately-contiguous target
and mask tensors suitable for the trainStepAdamF32Yolov1 FFI.
The per-record layout is target (perCell*gH*gW*4) || mask (gH*gW*4) || numBoxes (4) || raw_boxes (56*20); only the target and mask are extracted.
Returns (target_concat, mask_concat) sized batch * perCell*gH*gW*4 and
batch * gH*gW*4.
Pass the caller's real grid. This used to hardcode the Pets 7×7 record
(7200 bytes/record) while the caller sliced at dio.labelBytesPerRecord —
25428 at VisDrone-448/14×14. Reading at the wrong stride pairs each image
with a target lifted out of a different record, which nothing downstream can
see because the output shape is still correct. The FFI now rejects a stride
that disagrees with the buffer.
Bbox-aware horizontal flip for a YOLOv1 batch. Per-image p=0.5
coin (xorshift64 seeded by seed); when flipped, reverses image
along W, target along gridW, mask along gridW, and replaces the
x_cell channel with 1 - x_cell on cells where mask=1 (since the
cell itself mirrors). Returns the augmented (images, target, mask)
triple as fresh ByteArrays; inputs are not modified.
See planning/archive/yolo_final.md Phase 3. LEGACY — superseded by
yoloAugment (Phase 3b) which operates on raw bboxes.
Unified bbox-aware augmentation for YOLOv1: per-image hflip + random
crop, with target+mask re-encoded from the transformed raw bboxes
so the geometric correspondence is exact. Replaces yoloHflip for
Phase 3b once preprocessor stores raw bboxes alongside the
pre-encoded target.
images: f32 image batch[B, C, H, W]boxes: per-record YOLOv1 label block (target 5880 + mask 196 + numBoxes 4 + raw_boxes 1120 = 7200 bytes/record). Only the numBoxes + raw_boxes tail is read.hflipProb,cropProb: per-image Bernoulli probabilities.cropMinScale: crop side ∈[cropMinScale, 1.0] × imgW(paper's ±20% jitter → 0.8).seed: xorshift seed.
Returns (new_image, new_target, new_mask) as fresh ByteArrays.
See planning/archive/yolo_final.md Phase 3.
Photometric HSV jitter (YOLO-style) for a normalized image batch [B,C,H,W].
Three multiplicative gains 1 + U(-1,1)·gain on hue (mod 360°), saturation,
and value, one draw per image. Image-only — touches no labels, so it composes
with any detector target. imagenetNorm=1 de-norms to [0,1] sRGB, applies the
gains in HSV space, clamps, and re-norms (the FPN loader stores images
ImageNet-normalized). Returns a fresh image batch; the input is not modified.
Horizontal flip of an FPN image [B,C,H,W] + its flat [P3|P4|P5]
multi-scale target, one p=prob coin per image. A flip is shape-invariant, so
every GT keeps its scale AND best-shape anchor: the image and each scale's grid
mirror columns, and the in-cell x-offset tx becomes 1-tx on assigned cells
(obj=1). This matches encode_targets_fpn exactly — no re-encode from boxes is
needed (the FPN record stores none). scalesFlat is int32-LE pairs [g_s, A_s]
per scale; perAnchor=15 fixed. Returns (image', target') as fresh ByteArrays.
Box-aware affine (scale + translate) of an FPN image [B,C,H,W] + its flat
[P3|P4|P5] target, one p=prob coin per image.
Unlike fpnHflip, this is NOT shape-invariant: scaling changes max(w,h),
which changes both the FPN level a box lands on and which anchor wins the
wh-IoU. So the target is rebuilt from boxes — decoded out of the assigned
slots (every one of which is an exact encoding of one box), transformed,
clipped to the frame, filtered, and re-encoded by the same rules
encode_targets_fpn uses. Boxes already lost to a same-slot collision on
disk stay lost, which is correct: they are absent from the training target too.
⚠ The scale range is the load-bearing knob on this dataset. VisDrone objects
are 2–5 px after the 448 resize, so scaling down pushes them below P3's
stride-8 resolution; whThrPx and areaThr drop what the transform has
destroyed rather than encoding a degenerate target for it. Empty regions take
0.0, which in these ImageNet-normalized images is the dataset mean colour.
scalesFlat is int32-LE pairs [g_s, A_s]; anchorsFlat is f32 (w,h)
pairs packed per scale in the same order. Returns (image', target').
Convert a uint8 mask ByteArray (one byte per pixel) into a little-endian
int32 ByteArray of 4× the size. Pets loadPets returns masks as packed
uint8; trainStepAdamF32Seg expects int32 per-pixel class labels.
Paired horizontal flip for segmentation: flips the f32 image [B,C,H,W]
and the uint8 per-pixel mask [B,H,W] together, one coin per image, so the
pixel correspondence survives. A flip is a pure column permutation, so the
mask stays exact (no label interpolation). Returns (image', mask').
See lean_f32_seg_hflip_pair.
Per-batch segmentation confusion matrix. logits is f32 [B,NC,H,W],
masks is u8 [B,H,W] (per-pixel class). Returns int64 LE [NC*NC]
counts conf[true*NC + pred] (argmax over channels), for mIoU
accumulation across batches. planning/archive/unet_demo_v2.md Workstream A.
Convert little-endian int32 token IDs to f32, element for element.
Feeds the idsInput tokenPositionEmbed path: model input is [B, T]
f32 ids, one-hot built in-graph — the host-side [B, V·T] one-hot
buffer disappears. Exact for ids < 2²⁴.
Shuffle images and labels in-place (Fisher-Yates), applying the SAME permutation to both. Returns (shuffled images, shuffled labels).
labelBytes is the label's bytes per record — 4 for a classification
scalar, but a whole tensor for detection/segmentation (the FPN detector's
is 185220 floats = 740880 bytes). It used to be hardcoded to 4 in the FFI,
which permuted the images while leaving multi-float targets in place and so
destroyed the image/target pairing every epoch on every detector and
segmentation trainer. Pass dio.labelBytesPerRecord; never a literal.
Mixup (Zhang et al. 2017) — λ ~ Beta(α, α), x_mixed[i] =
λ·x[i] + (1-λ)·x[π(i)]. Returns the mixed image batch.
Pair with mixupSoftLabels using the SAME seed + alpha.
Soft labels for the mixup. Pair with mixupImages (same seed + alpha).
Output shape: [batch, nClasses] f32, with label smoothing applied.
CutMix (Yun et al. 2019) — paste a random rectangle from x[π(i)]
onto x[i]. Pair with cutmixSoftLabels (same seed + alpha).
KNN-Mixup — like Mixup but pair[i] is the nearest neighbor of i in
pixel-space L2 distance, not a random permutation. Mixes each sample
with its closest manifold sibling in the batch. Pair with
knnMixupSoftLabels using SAME images + seed + alpha.
KNN-Mixup soft labels. Needs the original images to recompute the
same KNN pairing the _images call used.
Random Erasing (Zhong et al. 2017) — with probability prob, fill a
random rectangle (relative area 2–33%, aspect 0.3–3.3) with N(0,1)
noise. Per-image independent. Labels unchanged.
RandAugment-Color (Cubuk et al. 2019, color-only subset). Per image,
apply nOps random ops drawn from {identity, brightness, contrast,
color, autocontrast} with magnitude m (0–10, paper default 9).
Geometric ops (rotate / shear / translate) are TODO.
imagenetNorm = 1 tells the kernel the incoming images are
ImageNet-mean/std normalized (Imagenette / Imagewoof); the kernel
de-normalizes to [0,1] sRGB, applies ops, then re-normalizes. Pass
0 for already-in-[0,1] datasets (CIFAR, MNIST).
Element-wise subtract: a − b. Used by SWAG for per-epoch deviation
snapshots p − swaMean.
Sample batch random sequences of length seqLen from a token stream.
Returns a single flat ByteArray of size 2 * batch * seqLen * 4
containing input IDs followed by next-token target IDs (both int32 LE).
GradCAM closed-form (Zhou 2016 CAM). For nets ending GAP+dense,
heat[i,j] = ReLU(Σ_k W[k, tgt] · A[k, i, j]), max-normalized to
[0, 1]. denseW is [C, NC] row-major, lastConv is [B, C, H, W]
NCHW. Returns [H, W] f32 for the chosen batchIdx.
Recompute logits from a pre-GAP activation. Returns [NC] f32 for
a single image (batchIdx). Used so the GradCAM exe can pick a class
via argmax without running the full forward a second time.
SWAG sample weights (Maddox et al. 2019). Given the SWA mean,
SWA-of-θ² (for diagonal variance), and the K most-recent per-epoch
deviation snapshots packed row-major as K × nParams f32, draw
one sample from the SWAG posterior N(μ, ½Σ_diag + ½ Σ_low).