| 952 | self.assertTrue(mx.allclose(grad, expected_grad)) |
| 953 | |
| 954 | def test_conv_groups_grad(self): |
| 955 | def fn(x, w): |
| 956 | num_groups = x.shape[-1] // w.shape[-1] |
| 957 | return mx.conv1d(x, w, groups=num_groups) |
| 958 | |
| 959 | def fn_gt(x, w): |
| 960 | num_groups = x.shape[-1] // w.shape[-1] |
| 961 | group_size = w.shape[-1] |
| 962 | ws = w.reshape(num_groups, -1, *w.shape[1:]).split(num_groups) |
| 963 | xs = x.reshape(*x.shape[:-1], num_groups, -1).split(num_groups, axis=-2) |
| 964 | return mx.concatenate( |
| 965 | [mx.conv_general(x.squeeze(-2), w.squeeze(0)) for x, w in zip(xs, ws)], |
| 966 | axis=-1, |
| 967 | ) |
| 968 | |
| 969 | mx.random.seed(3) |
| 970 | |
| 971 | w = mx.random.normal(shape=(2, 3, 1)) |
| 972 | x = mx.random.normal(shape=(1, 5, 2)) |
| 973 | cotans = (mx.ones(shape=(1, 3, 2)),) |
| 974 | grads = mx.vjp(fn, (x, w), cotans)[1] |
| 975 | expected = mx.vjp(fn_gt, (x, w), cotans)[1] |
| 976 | self.assertTrue(mx.allclose(expected[0], grads[0])) |
| 977 | self.assertTrue(mx.allclose(expected[1], grads[1])) |
| 978 | |
| 979 | w = mx.random.normal(shape=(2, 3, 2)) |
| 980 | x = mx.random.normal(shape=(1, 5, 4)) |
| 981 | cotans = (mx.ones(shape=(1, 3, 2)),) |
| 982 | grads = mx.vjp(fn, (x, w), cotans)[1] |
| 983 | expected = mx.vjp(fn_gt, (x, w), cotans)[1] |
| 984 | self.assertTrue(mx.allclose(expected[0], grads[0])) |
| 985 | self.assertTrue(mx.allclose(expected[1], grads[1])) |
| 986 | |
| 987 | w = mx.random.normal(shape=(6, 3, 2)) |
| 988 | x = mx.random.normal(shape=(1, 5, 4)) |
| 989 | cotans = (mx.ones(shape=(1, 3, 6)),) |
| 990 | grads = mx.vjp(fn, (x, w), cotans)[1] |
| 991 | expected = mx.vjp(fn_gt, (x, w), cotans)[1] |
| 992 | self.assertTrue(mx.allclose(expected[0], grads[0])) |
| 993 | self.assertTrue(mx.allclose(expected[1], grads[1])) |
| 994 | |
| 995 | # Test 2D |
| 996 | w = mx.random.normal(shape=(2, 3, 3, 1)) |
| 997 | x = mx.random.normal(shape=(1, 5, 5, 2)) |
| 998 | cotans = (mx.ones(shape=(1, 3, 3, 2)),) |
| 999 | grads = mx.vjp(fn, (x, w), cotans)[1] |
| 1000 | expected = mx.vjp(fn_gt, (x, w), cotans)[1] |
| 1001 | self.assertTrue(mx.allclose(expected[0], grads[0])) |
| 1002 | self.assertTrue(mx.allclose(expected[1], grads[1])) |
| 1003 | |
| 1004 | # Test with flip |
| 1005 | def fn(x, w): |
| 1006 | num_groups = x.shape[-1] // w.shape[-1] |
| 1007 | return mx.conv_general(x, w, groups=num_groups, flip=True) |
| 1008 | |
| 1009 | def fn_gt(x, w): |
| 1010 | num_groups = x.shape[-1] // w.shape[-1] |
| 1011 | group_size = w.shape[-1] |