(self)
| 406 | self.assertLess((y_q - y_hat).abs().max(), 1e-3) |
| 407 | |
| 408 | def test_qvm(self): |
| 409 | key = mx.random.key(0) |
| 410 | k1, k2 = mx.random.split(key) |
| 411 | tests = product( |
| 412 | [128, 64, 32], # group_size |
| 413 | [2, 3, 4, 5, 6, 8], # bits |
| 414 | [32, 128, 256], # M |
| 415 | [128, 256, 67], # N |
| 416 | [0, 1, 3, 8], # B |
| 417 | ) |
| 418 | for group_size, bits, M, N, B in tests: |
| 419 | with self.subTest(shape=(B, M, N), group_size=group_size, bits=bits): |
| 420 | if M < group_size: |
| 421 | continue |
| 422 | x_shape = (1, N) if B == 0 else (B, 1, N) |
| 423 | w_shape = (N, M) if B == 0 else (B, N, M) |
| 424 | x = mx.random.normal(shape=x_shape, key=k1) |
| 425 | w = mx.random.normal(shape=w_shape, key=k2) |
| 426 | w_q, scales, biases = mx.quantize(w, group_size, bits) |
| 427 | w_hat = mx.dequantize(w_q, scales, biases, group_size, bits) |
| 428 | y_q = mx.quantized_matmul( |
| 429 | x, w_q, scales, biases, False, group_size, bits |
| 430 | ) |
| 431 | y_hat = x @ w_hat |
| 432 | self.assertEqual(y_q.shape, y_hat.shape) |
| 433 | self.assertLess((y_q - y_hat).abs().max(), 1e-3) |
| 434 | |
| 435 | def test_qvm_splitk(self): |
| 436 | key = mx.random.key(0) |
nothing calls this directly
no test coverage detected