Apple AxLearn

Quantization-test and evaluation fixes for Apple's JAX research framework.

Contributed to Apple’s JAX-based research framework by improving AQT quantization test coverage and resolving evaluation issues that strengthened low-precision training workflows.