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:…