(self)
| 9 | |
| 10 | class TestQuantized(mlx_tests.MLXTestCase): |
| 11 | def test_quantize_dequantize(self): |
| 12 | w = mx.random.normal(shape=(128, 512)) |
| 13 | for gs in [32, 64, 128]: |
| 14 | for b in [2, 3, 5, 6, 4, 8]: |
| 15 | with self.subTest(gs=gs, b=b): |
| 16 | w_q, scales, biases = mx.quantize(w, group_size=gs, bits=b) |
| 17 | w_hat = mx.dequantize(w_q, scales, biases, gs, b) |
| 18 | errors = (w - w_hat).abs().reshape(*scales.shape, -1) |
| 19 | eps = 1e-6 |
| 20 | self.assertTrue((errors <= (scales[..., None] + eps).abs()).all()) |
| 21 | |
| 22 | # test quantize/dequantize 0s |
| 23 | a = mx.zeros((256, 512)) |
| 24 | for gs in [32, 64, 128]: |
| 25 | for b in [2, 3, 4, 5, 6, 8]: |
| 26 | w_q, scales, biases = mx.quantize(a, gs, b) |
| 27 | a_hat = mx.dequantize(w_q, scales, biases, gs, b) |
| 28 | self.assertTrue(mx.all(a_hat == 0)) |
| 29 | |
| 30 | def test_mxfp4_quantize_dequantize(self): |
| 31 | lut = mx.array( |
nothing calls this directly
no test coverage detected