| 273 | self.assertTrue(mx.allclose(vjp_out[0], expected_out)) |
| 274 | |
| 275 | def test_qmm_jvp(self): |
| 276 | key = mx.random.key(0) |
| 277 | k1, k2 = mx.random.split(key) |
| 278 | |
| 279 | bits = 8 |
| 280 | group_size = 64 |
| 281 | M = 64 |
| 282 | N = 128 |
| 283 | K = 128 |
| 284 | |
| 285 | x = mx.random.normal(shape=(2, M, K), key=k1) |
| 286 | x_tan = mx.ones(shape=(2, M, N)) |
| 287 | |
| 288 | transposes = [True, False] |
| 289 | for transposed in transposes: |
| 290 | w = mx.random.normal(shape=(N, K) if transposed else (K, N), key=k2) |
| 291 | w_q, scales, biases = mx.quantize(w, group_size, bits) |
| 292 | |
| 293 | def fn(x): |
| 294 | return mx.quantized_matmul( |
| 295 | x, w_q, scales, biases, transposed, group_size, bits |
| 296 | ) |
| 297 | |
| 298 | _, jvp_out = mx.jvp(fn, primals=(x,), tangents=(x_tan,)) |
| 299 | |
| 300 | expected_out = mx.quantized_matmul( |
| 301 | x_tan, w_q, scales, biases, transposed, group_size, bits |
| 302 | ) |
| 303 | self.assertTrue(mx.allclose(jvp_out[0], expected_out)) |
| 304 | |
| 305 | def test_qmm_shapes(self): |
| 306 | key = mx.random.key(0) |