| 3057 | np.testing.assert_allclose(out, out_t, atol=1e-4) |
| 3058 | |
| 3059 | def test_roll(self): |
| 3060 | x = mx.arange(10).reshape(2, 5) |
| 3061 | |
| 3062 | for s in [-2, -1, 0, 1, 2]: |
| 3063 | y1 = np.roll(x, s) |
| 3064 | y2 = mx.roll(x, s) |
| 3065 | self.assertTrue(mx.array_equal(y1, y2).item()) |
| 3066 | |
| 3067 | y1 = np.roll(x, (s, s, s)) |
| 3068 | y2 = mx.roll(x, (s, s, s)) |
| 3069 | self.assertTrue(mx.array_equal(y1, y2).item()) |
| 3070 | |
| 3071 | shifts = [ |
| 3072 | 1, |
| 3073 | 2, |
| 3074 | -1, |
| 3075 | -2, |
| 3076 | (1, 1), |
| 3077 | (-1, 2), |
| 3078 | (33, 33), |
| 3079 | ] |
| 3080 | axes = [ |
| 3081 | 0, |
| 3082 | 1, |
| 3083 | (1, 0), |
| 3084 | (0, 1), |
| 3085 | (0, 0), |
| 3086 | (1, 1), |
| 3087 | ] |
| 3088 | for s, a in product(shifts, axes): |
| 3089 | y1 = np.roll(x, s, a) |
| 3090 | y2 = mx.roll(x, s, a) |
| 3091 | self.assertTrue(mx.array_equal(y1, y2).item()) |
| 3092 | |
| 3093 | def test_roll_errors(self): |
| 3094 | x = mx.array([]) |