Enable AQT quantization test

Author: RissyRan

MaxText/tests/quantizations_test.py

@@ -105,7 +105,6 @@ def test_configure_quantization_is_int8(self):
       quant = _configure_quantization(quant_str="int8", mode_str=quant_mode)
       self.assertNotEqual(quant, None)
 
-  @pytest.mark.skip(reason="b/400476456 Tests are currently flaking / failing due to JAX 0.5.1 upgrade")
   def test_aqt_quantization(self):
     # Without quantization
     inputs, res_einsum, res_dg = _apply()