(self)
| 585 | mx.eval(y) |
| 586 | |
| 587 | def test_small_matrix(self): |
| 588 | for w_shape in [(8, 256), (1, 8, 256), (3, 8, 256)]: |
| 589 | with self.subTest(w_shape=w_shape): |
| 590 | w = mx.random.normal(shape=(w_shape)) |
| 591 | w_q, scales, biases = mx.quantize(w) |
| 592 | w_hat = mx.dequantize(w_q, scales, biases) |
| 593 | |
| 594 | # Test qmv |
| 595 | for shape in [(3, 1, 256), (3, 4, 256)]: |
| 596 | x = mx.random.normal(shape=shape) |
| 597 | y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=True) |
| 598 | y_hat = x @ mx.swapaxes(w_hat, -1, -2) |
| 599 | self.assertEqual(y_q.shape, y_hat.shape) |
| 600 | self.assertLess((y_q - y_hat).abs().max(), 1e-3) |
| 601 | |
| 602 | # Test qmm_t |
| 603 | x = mx.random.normal(shape=(3, 10, 256)) |
| 604 | y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=True) |
| 605 | y_hat = x @ mx.swapaxes(w_hat, -1, -2) |
| 606 | self.assertEqual(y_q.shape, y_hat.shape) |
| 607 | self.assertLess((y_q - y_hat).abs().max(), 1e-3) |
| 608 | |
| 609 | # Test qvm |
| 610 | x = mx.random.normal(shape=(3, 1, 8)) |
| 611 | y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=False) |
| 612 | y_hat = x @ w_hat |
| 613 | self.assertEqual(y_q.shape, y_hat.shape) |
| 614 | self.assertLess((y_q - y_hat).abs().max(), 1e-3) |
| 615 | |
| 616 | # Test qmm |
| 617 | x = mx.random.normal(shape=(3, 10, 8)) |
| 618 | y_q = mx.quantized_matmul(x, w_q, scales, biases, transpose=False) |
| 619 | y_hat = x @ w_hat |
| 620 | self.assertEqual(y_q.shape, y_hat.shape) |
| 621 | self.assertLess((y_q - y_hat).abs().max(), 1e-3) |
| 622 | |
| 623 | def test_non_multiples(self): |
| 624 | w = mx.random.normal(shape=(33, 256)) |
nothing calls this directly
no test coverage detected