We use the Neural Tangents (NT) library [1]}
built on top of JAX [2]}
for all CIFAR10 experiments, and the ImageNet codebase from FLAXhttps://github.com/google/flax/blob/main/examples/imagenet/README.md[3]} for ResNets and Mixer experiments on ImageNet.
