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.
## References
- He et al. 2015, *Delving Deep into Rectifiers* (He init). <https://arxiv.org/abs/1502.01852>
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.
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 — it avoids millions of
Lean-level push calls; see dropoutMask below.
recs[idx[0]] ++ recs[idx[1]] ++ … in one C pass (idx = u32 LE); every record
is recBytes. A replay buffer's batch, e.g. pong-dqn's f32 state records. dst is
reused in place when it is unshared and the right size (last batch's buffer).
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 step of hundreds of milliseconds — 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 — 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 difference is the whole distinction between the two regularisers. The reference draws
bernoulli(key, keep, x.shape) for the classifier (the classifier dropout in the JAX emitter,
jax/Jax/Codegen.lean) against (branch.shape[0],) + (1,)*(ndim-1) for stochastic depth. 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. The shapes agree and the
emitted graph is identical; only the distribution differs. Proofs.dropPath_scales_uniformly
states what would be wrong; on this side it is the loop bound, and no gate checks 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)) so a net running both
regularisers does not correlate them; the caller adds that offset. Sharing a seed 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. The draw loop itself is dropoutFill (C, ffi/f32_helpers.c): at a global batch of
256 and width 1280, n = 327,680 per step, and a Lean-side loop at that size costs a large
share of the step. The guards and the seeding argument stay here.
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
The [θ | m | v] prefix of a packed Adam train step's output, nP floats each: the state
the next step takes back as (θ.append m).append v.
Equations
Instances For
Extract the loss (last float32) from a train_step output.
Equations
- F32.extractLoss out lossIdx = F32.read out lossIdx.toUSize
Instances For
Decode the little-endian int32 label at record i of a packed label buffer.
Reads all four bytes. Reading byte 0 alone (lbl.get! (4 * i)) is label % 256: invisible
on the 10-class datasets, but on 1000-class ImageNet most labels exceed 255, and a correct
prediction could then match only on classes 0..255.
Equations
- F32.readLabel lbl i = readU32LE lbl (4 * i)
Instances For
Argmax over n float32 values starting at element offset off.
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. A literal 10 would be right on the 10-class datasets
and would confine every 1000-class prediction to labels 0..9.
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. 240*240 = 57600
for BraTS). Zero-copy.
Equations
Instances For
Gather count records of an Imagenette-format file held as RAW bytes (IO.FS.readBinFile)
by index — idx is count little-endian u32 record indices — into an ImageNet-normalised
f32 batch [count, 3, imgSize, imgSize], exactly as loadImagenetteSized would have
converted them. The u8 records stay resident (196 KB per 256² image against 786 KB as f32),
which is what lets a 43k-image PlantVillage part, or the union of two, sit in memory.
Gather count chips of a flat f32 [n, C, S, S] set by index (idx is count
little-endian u32 chip indices), each under one of the eight symmetries of the square
drawn from seed and its batch position (bit 0 flips columns, bit 1 rows, bit 2
transposes). The remote-sensing demo's training augmentation: a nadir chip has no up,
so every one of the eight is a label-preserving view.
Load Imagenette keeping pixels RAW uint8 (1 byte/element, channel-first), labels i32 — one
byte per element where loadImagenetteSized holds four, which is what makes Imagenette fit
a Jetson's unified memory (deploy/ORIN.md §4). Batches are normalised one at a time
by sliceU8NormPad. Opt-in via LEAN_MLIR_IMAGENETTE_U8=1 (see loadData).
One batch of raw-uint8 images — count records of pixelsPerImage bytes starting at
start, in a dataset of total records — as normalised NCHW f32, zero-padded past total
exactly like sliceImagesPad. The normalisation is bit-identical to loadImagenetteSized.
Imagenette LABELS only (i32 class ids) and the record count, without reading any pixels — the
streamed train path (LEAN_MLIR_IMAGENETTE_STREAM=1). Rejects a file whose size is not exactly
count records at imgSize.
[0, 1, …, n-1] as little-endian u32: the unshuffled index array for readU8NormPad. Shuffle it
with shuffleU8 indices labels n 4 4 seed, which applies the same permutation shuffleU8 applies
to a resident image buffer.
sliceU8NormPad without the resident buffer: slot r of the batch is record
indices[start + r] (little-endian u32), read from the Imagenette file at path and normalised by
the same code. With indices shuffled alongside the labels, the batch is byte-identical to
slicing the equally-shuffled resident buffer.
Load a BraTS (MSD Task01_BrainTumour) binary file at the given in-plane
size and channel count. Returns (images f32 ByteArray, masks uint8 ByteArray, count).
Images are imgSize×imgSize×channels, channel-first: the 4 modalities
(FLAIR / T1w / T1gd / T2w) of a 2D build, or 4·(2K+1) for a 2.5D build with
K context slices each side (preprocess_brats.py --context K), the centre
slice's modalities at channels 4K..4K+3.
Unlike the RGB datasets these carry no ImageNet normalization: the loader
inverts the uint8 quantization scripts/datasets/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).
The record size is not in the file's header, so the loader checks the file
length against count · (channels·S² + S²) and refuses a mismatch.
YOLOv1 detection-bin loader (target+mask format, 224 input / 7×7 grid). 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), matching the detection
dataset's labelBytesPerRecord in Train. The Lean dispatcher (runTraining) splits this into target + mask before
calling trainStepAdamF32Yolov1. See scripts/datasets/preprocess_visdrone.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 path.
Anchor-format detection loader. 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. 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: the stride is derived from it (7200 bytes/record at 224/7×7, 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 rejects a stride that disagrees with the buffer.
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. Reads records that store the 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.
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 knob that matters most 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').
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.
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). A literal 4 would permute the images while
leaving multi-float targets in place, destroying the image/target pairing
every epoch. Pass dio.labelBytesPerRecord, never a literal.
shuffle for raw-uint8 image buffers (1 byte per element instead of 4). Same RNG and swap
order, so for the same n and seed it applies the same permutation as shuffle.
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).