| 699 | self.assertLess(mx.abs(gb1 - gb2).max() / mx.abs(gb1).mean(), 1e-5) |
| 700 | |
| 701 | def test_fast_transforms(self): |
| 702 | x = mx.random.uniform(shape=(2, 2, 8)) |
| 703 | |
| 704 | defaults = (8, False, 10000.0, 1.0, 0) |
| 705 | dims, traditional, base, scale, offset = defaults |
| 706 | |
| 707 | # VJP |
| 708 | _, vjp_out = mx.vjp(lambda x: rope_orig(x, *defaults), (x,), (mx.ones_like(x),)) |
| 709 | _, vjp_fast_out = mx.vjp( |
| 710 | lambda x: mx.fast.rope( |
| 711 | x, dims, traditional=traditional, base=base, scale=scale, offset=offset |
| 712 | ), |
| 713 | (x,), |
| 714 | (mx.ones_like(x),), |
| 715 | ) |
| 716 | self.assertTrue(mx.allclose(vjp_out[0], vjp_fast_out[0])) |
| 717 | |
| 718 | # JVP |
| 719 | _, jvp_out = mx.jvp(lambda x: rope_orig(x, *defaults), (x,), (mx.ones_like(x),)) |
| 720 | _, jvp_fast_out = mx.jvp( |
| 721 | lambda x: mx.fast.rope( |
| 722 | x, dims, traditional=traditional, base=base, scale=scale, offset=offset |
| 723 | ), |
| 724 | (x,), |
| 725 | (mx.ones_like(x),), |
| 726 | ) |
| 727 | self.assertTrue(mx.allclose(jvp_out[0], jvp_fast_out[0])) |
| 728 | |
| 729 | # VMAP |
| 730 | x = mx.random.uniform(shape=(2, 2, 2, 8)) |
| 731 | vmap_out = mx.vmap(lambda x: rope_orig(x, *defaults))(x) |
| 732 | vmap_fast_out = mx.vmap( |
| 733 | lambda x: mx.fast.rope( |
| 734 | x, dims, traditional=traditional, base=base, scale=scale, offset=offset |
| 735 | ) |
| 736 | )(x) |
| 737 | self.assertTrue(mx.allclose(vmap_out, vmap_fast_out)) |
| 738 | |
| 739 | @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") |
| 740 | def test_custom_kernel_basic(self): |