JAX
Runnable JAX and Flax NNX implementations: contrastive losses, linear attention, Q/K attention, and prototype readouts, with the math running live.
-
Measuring Attention's Geometry in JAX/Flax NNX
A runnable companion to the geometry-of-attention post: the telemetry that exports a trained head's actual query and key vectors, the offline replay that recomputes every row's winner under rescaled queries, the linear program that asks whether any key sits inside the others' convex hull, and the census of who owns whom through training. Plus the toy laws animated from their real formulas: territories re-carving, a body crossing the hull, the far sky losing its sign.
-
The Geometry of Attention Is a Choice of Kernel
An attention head is only ever evaluated at the tokens of the sequence, but its formula accepts any query vector at all. So between the tokens there is a whole continuous space, and the head silently carves it into territories: for every possible query, some key takes the top weight. This post draws that map for the two attention laws this series has trained head to head. The dot-product law slices space into infinite wedges of sky that all meet at the origin, a gravity that reads only bearings, where a query's length is a temperature and a key inside the others' hull can never win. The Yat kernel carves bounded neighborhoods around bodies that always own their ground, a gravity that reads places, where length is an address. One theorem per map, one live map per theorem; the trained transformers appear only in a coda, to confirm they were never free to draw anything else.
-
Softmax-Free Attention in JAX/Flax NNX
A runnable companion to the compatibility-kernel post: the attention module where one branch computes softmax and the other computes the Yat kernel with no exponential anywhere, the parameter-matched training harness, the telemetry that measured our bounded-scores belief dead, and the checkpointed attention maps. Every number is from the real Kaggle runs.
-
The Kernel Between the Roles
The QK post in this series ended on a construction it refused to build: keep the query and key roles, but replace the bilinear-then-exponential score with a genuine Mercer kernel between them. This post builds it, with the kernel's full form: a per-head learned bias inside the square, the term the universality theorem requires, and a per-head learned softening. Because the kernel is nonnegative by construction, attention needs no softmax at all and routing becomes a literal Nadaraya-Watson smoother. Trained head to head at matched parameters and per-variant swept learning rates, the kernel transformer lands within 1.1 percent of softmax on character-level Shakespeare, and the differences that survive are the interesting part: no gauge, no max-trick, a mass channel softmax cannot represent, and two of our own assumptions measured dead.
-
An Error Controller for a Trained Net, in JAX
A runnable companion to the depth-on-demand post: the leapfrog classifier trained with lax.scan at fixed depth, the step-doubling controller that re-renders it to tolerance at inference, the honest work accounting (probes included), and the measured tol^(-1/3) power law. Every figure is rendered from the real Kaggle run.
-
Depth on Demand
The last post made depth a resolution: layers are time steps of a learned flow, and running more of them just renders the same trajectory finer. But every camera knows not to spend equal film on empty sky. This post gives a trained network the integrator's next tool, an error controller that chooses its own step size per input, with no retraining: the same weights, rendered to tolerance. The controller reproduces the reference verdicts at a fraction of the steps, its cost follows the integrator's textbook one-third power law, and the map of where it spends is a genuine surprise: effort tracks the stiffness of the learned flow, not the difficulty of the classification.
-
Reversible Backprop as a custom_vjp in JAX
A runnable companion to the memory post: the momentum block and its exact inverse, the custom_vjp whose backward pass reconstructs the trajectory instead of storing it, XLA's memory_analysis as the measuring instrument, and the (1/mu)^L float cliff reproduced in numpy float32. Every number is from the real Kaggle run.
-
Backprop Without the Memory
Training memory is a tax nobody chose: backprop must hold every activation of the forward pass hostage until the backward pass consumes it, so depth costs memory even when it costs little compute. This post spends the invertibility this series proved two posts ago: a momentum residual block can be run backward, so the backward pass can recompute the past instead of storing it. Measured on the same network, standard backprop's activation memory grows from 13 MB to 674 MB as depth goes 8 to 512; the reversible pass holds flat at 3.2 MB, pays 24% in step time, and returns the same gradient, until a one-line arithmetic of friction and float noise says it cannot.
-
Building the Energy-Conserving Net in JAX/Flax NNX
A runnable companion to the energy-conservation post: the HNN pendulum field as the symplectic gradient of one learned scalar, the plain field model it beats on drift, and the leapfrog classifier whose residual block is a kick-drift-kick step of a learned potential, all as Flax NNX modules with lax.scan doing depth. Every figure is rendered from the real Kaggle run.
-
A Network That Conserves Energy
A pendulum never forgets its energy, and a trained network has no such number to forget. This post builds a residual network whose hidden state carries a conservation law by construction: the block is a symplectic step of a learned energy, so the quantity is held by the architecture, not encouraged by a loss. The learned pendulum keeps its energy to 0.6% where a plain field model leaks 36%, the classifier lands in the pack on accuracy, and the law pays where composition fails: trained at depth 16 and run at four times that, the plain net gives up 31 points on spirals while the leapfrog net holds.
-
Running the Survival Trial, in JAX/Flax NNX
A runnable companion to the survival-model trial: the Yat DeepSurv trunk in Flax NNX, the Cox partial-likelihood loss, an LR-fair per-model training loop with best-epoch selection, the concordance / integrated-Brier / time-dependent-AUC evaluation, and the classical baselines wired through sksurv and lifelines. Every figure is a real number from a real run across five datasets, with the prototypes and the risk stratification captured as they form over training.
-
A Velocity Ledger for Transformers, in JAX/Flax NNX
A runnable companion: the pre-norm Transformer block as a forward-Euler step, then the residual-stream velocity ledger as one line of Flax NNX state (mu = 0 recovers plain), the ngpt-lite retraction variant, best-val early-stopped training, and the depth telemetry (path length and turning angle per sub-update). Four parameter-matched char-level GPTs that tie on quality and split on dynamics: the ledger's residual-stream path is a third as long and half as sharp.
-
Solving It and Descending It, in JAX/Flax NNX
A runnable companion to the solve-vs-descend post: the Yat kernel and its Gram matrix, the exact kernel ridge solve via Cholesky, the same kernel as a Flax NNX module trained by AdamW with LR sweeps and best-epoch selection, the measured timing wall, minibatching through 511k rows, and the conv trunk the solve can never train. Every number is from the real Kaggle runs.
-
A White-Box DeepSurv, in JAX/Flax NNX
A runnable companion: the Cox partial-likelihood loss, a standard DeepSurv, and a Yat-kernel DeepSurv whose log-risk decomposes into prototype patients, all in Flax NNX. Concordance evaluation, exact convex attribution, cohort deletion, and OOD abstention as short array operations. Every number is from a real run on METABRIC breast-cancer survival.
-
Calibrating a Bounded Net, in JAX/Flax NNX
A runnable companion: build the matched Yat and ReLU MLPs in Flax NNX with the same softmax head, then measure their honesty. The reliability diagram and ECE, temperature scaling fit on a held-out split, NLL and Brier, and the two out-of-distribution channels, kernel-field magnitude versus softmax confidence, all in JAX with every number from a real three-seed run.
-
Building the Second Layer by Hand, in JAX/Flax NNX
A runnable companion: build a whole second feature layer by hand in JAX, on top of the hand-built first. Named min-AND combinations of layer-1 edges (junctions, continuations, bends, stripes) feed the same constructed Yat head, no training anywhere. It reproduces the flat rung: 83.3% at layer 1, 82.9% with both, 78.8% from relations alone, and counts the combinatorial wall of 224 pairwise and 4,630 three-way types where construction stops.
-
Distillation as Kernel Transfer, in JAX/Flax NNX
A runnable companion: the five-run distillation experiment in JAX/Flax NNX. Train a teacher CNN, extract its class-similarity kernel S = E[softmax(z/T) softmax(z/T)ᵀ], train a student on nothing but pairwise relations (no labels, no soft targets), and measure it against the label ceiling and the random floor with a linear and a nearest-centroid probe. Every number is from a real run, with six GIFs that animate the kernel assembling, the temperature dial, the handoff, the spectrum inheritance, the probe race, and the inherited mistakes.
-
Editing a Deep Equilibrium Network, in JAX/Flax NNX
A runnable companion: build the weight-tied Yat equilibrium operator in Flax NNX, then teach a class by appending rows to the readout (F untouched, exact) or into the shared dynamics (one paste, present at every depth), measure the contraction certificate with power iteration and bisect one gain to restore it, audit the drift of 520 old fixed points, watch a layer-only edit evaporate, and forget by masking. Every number is from a real run.
-
Skip Connections With Inertia, in JAX/Flax NNX
A runnable companion: the residual block as a forward-Euler step, then the momentum residual network as a Flax NNX module with one extra state, a velocity the blocks write into. Train both on the rings task a first-order flow cannot separate exactly, watch the training crystallize, and run the trained network exactly backward until floating point, amplified by 1/mu per layer, steals the past.
-
Distillation Is a Geometry, Not an Answer Key
Knowledge distillation has a standing puzzle: Hinton's student recognized 98.6% of the digit 3s in the test set after training on a transfer set with every 3 deleted. An answer key cannot do that, so what actually crosses the wire? This post gives dark knowledge a data type, a class-similarity kernel, and runs the experiment that isolates it: a student trained on nothing but pairwise relations, no labels, no soft targets, no class names, measured against the label-trained ceiling and the random floor. With live experiments: watch the kernel accumulate from single outputs, turn the temperature knob on how much geometry leaks, train a relational student in the page, and watch whose spectrum the student grows into.
-
The Price List, in JAX/Flax NNX
A runnable companion to the price-list post: kernel ridge in JAX, the representer solve (K + lambda I) alpha = y, the RKHS-norm bill alpha^T K alpha, the effective dimension d_eff = sum lambda_k/(lambda_k + lambda) from the Gram spectrum, and a generalization sweep that draws the U-curve. Every number and every figure is from one analytic solve, no gradient descent.
-
A Network That Is a Fixed Point, in JAX/Flax NNX
A runnable companion: build the Yat deep-equilibrium network in JAX/Flax NNX. One shared operator F(z;x)=tanh(A·φ_W(z)+Ux+z0), solved to its fixed point by damped iteration, trained not by backprop-through-iterations but by implicit differentiation, the adjoint (I−Jᵀ)u=∂L/∂z* run by the same contraction. Plus the weight-tied maze operator that extrapolates from 11×11 to 27×27 by iterating longer. Every number is from a real run: 98.2% on two moons from 1700 shared parameters, ‖J‖ 0.66/0.92, 99.5% on mazes far larger than training.
-
A White-Box FFN: the Representer Theorem in JAX/Flax NNX
A runnable companion: build a transformer whose feed-forward block is a Yat kernel, so the FFN is exactly a representer sum over learned key-value memory slots. Train it on tinyshakespeare, then do four things you cannot do to an opaque ReLU FFN: read each memory slot, attribute an output to the slots that wrote it, edit one slot and watch generation change, and read off when the memory is out of its depth.
-
What a Weight Can Be, in JAX/Flax NNX
A runnable companion: compute the price list of a kernel in JAX. The eigenvalues are the kernel's spectral density, found with an FFT; the RKHS norm of a weight is a sum over them. A corner is affordable only under a Sobolev kernel, and the same numbers place the Yat kernel: universal and smooth, roomier than a Gaussian but not a Sobolev space.
-
Where a Weight Lives, in JAX/Flax NNX
A runnable companion: build the representer-theorem weight in JAX. A positive-definite kernel, the Gram matrix, a single linear solve for the coefficients, and the weight comes out as a combination of the data, f = sum alpha_i k(x_i, .). A linear weight cannot separate nested rings; the placed kernel weight does, read purely through the kernel as a similarity-weighted vote of the data.
-
The Hand-Built Network, in JAX/Flax NNX
A runnable companion: build the training-free image classifier from the post in JAX. The feature extractor is pure JAX (Sobel gradients, orientation binning, patch pooling); the classifier is a Flax NNX module holding k-means prototypes that votes with the Yat kernel. Nothing is trained, and it reproduces the 83.3% on Fashion-MNIST.
-
Constructing the Head on Learned Features, in JAX/Flax NNX
A runnable companion: train a small conv backbone in Flax NNX, then on its frozen features build a constructed Yat head with no gradient steps and compare. The constructed head lands within a couple of points of the trained one, and even a random backbone's features sort at 73% while its trained head is at chance.
-
Editing a Network by Hand, in JAX/Flax NNX
A runnable companion: build the prototype Yat-MLP in Flax NNX, then add a class by concatenating a few prototype rows and forget a class by masking them out, with no gradient steps. Class-incremental learning that matches a from-scratch build, and exact machine unlearning, both as array edits you can read. Every number is from a real run on Fashion-MNIST.
-
Your Neuron Is a Picture, in JAX/Flax NNX
A runnable companion: build the prototype MLP from the post in Flax NNX, train it on Fashion-MNIST, and watch the neurons. Pull the prototypes out as images, read a prediction as a vote over pictures, see the model abstain on out-of-distribution digits, check that random-init prototypes classify but stay noise, and track the prototypes migrating through a UMAP fit on the dataset as they train.
-
The Yat-Kernel MLP in JAX/Flax NNX
A runnable companion to What a Finite Kernel Buys an MLP: build a layer whose unit is the Yat kernel instead of a linear map plus an activation, assert it is positive definite and nonnegative, write down its exact finite feature map, train it end-to-end on two moons with no activation function, and measure the lazy-loading sparsity, the bounded off-distribution response, the RKHS capacity, and the force field that pulls each prototype onto its data.
-
The Three States of Information, in JAX
A runnable companion to The Three States of Information: train tiny models in JAX and measure the three states directly: the feature-covariance spectrum collapsing from high-rank (random) to a C−1-mode frame (structured), the distributional simplicity bias that fits low-order structure first (organized), the neural-collapse simplex where class-mean cosines lock onto −1/(C−1), and the alignment/uniformity split of contrastive learning running on two separate clocks. Four live JAX visualizations, every number an eigenvalue or a loss.
-
Latent on the Spectrum, in JAX
A runnable companion to Latent on the Spectrum: build a codebook as the spectral embedding of a label kernel in JAX (classical MDS with square-root eigenvalue scaling), watch a flat spectrum become the simplex and a graded one become the horseshoe, measure kernel-target alignment, split a representation into its between-class prototype frame and within-class information spectrum, and watch neural collapse grind the information to zero.
-
Q and K Projections in JAX/Flax NNX
A runnable companion to Why Attention Needs Q and K Projections: build scaled dot-product attention with separate query and key projections in Flax NNX, pull the bilinear form B = W_Q W_Kᵀ out of the module, split it into a symmetric metric and an antisymmetric directed part, wire a toy induction head, add RoPE, and measure the low-rank budget and the gauge freedom, all in plain JAX.
-
The Prototype Readout in JAX/Flax NNX
A runnable companion to The Readout is a Convex Combination of Prototypes: read the columns of W_out as output prototypes in Flax NNX, measure the convex/conic/affine/linear regimes numerically, then build a Nadaraya–Watson kernel readout that is convex by construction (nonnegative weights that sum to one, a point that never leaves the prototype hull), with the nonnegativity-vs-positive-definiteness distinction checked in code.
-
Auditing Latent Space Geometry in JAX
A runnable companion to the Welch-bound latent-space post: generate GIFs and implement the JAX metrics that tell you whether embeddings are collapsing, wasting rank, forming a simplex, or pressing against the Welch floor.
-
Cheap Attention in JAX/Flax NNX
A runnable companion to Cheap Attention: implement positive-feature linear attention in JAX and Flax NNX, watch the all-pairs ledger turn into a shared feature state, and see exactly where the N×N matrix disappears.
-
Organizing Randomness: Contrastive Learning in JAX
A block-by-block JAX + Optax implementation of six contrastive losses, each watched as a real animated GIF turning random 2D points into organized embeddings. The runnable companion to "Untangling the Moons."
-
Self-Attention as Kernel Regression in JAX/Flax NNX
A runnable companion to Attention is Explainable Because it is a Kernel: build scaled dot-product attention from scratch in Flax NNX, prove in code that it is exactly a Nadaraya–Watson kernel smoother, watch the separate q/k projections break positive-definiteness numerically, swap the exp-dot-product kernel for Gaussian, Yat, and linear kernels to see which keep the weights a convex partition of unity, read the temperature as a kernel bandwidth, and train a single head end-to-end to route to a marked token.