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.