| 68 | self.assertTrue(mx.array_equal(y, op(x), equal_nan=True)) |
| 69 | |
| 70 | def test_binary(self): |
| 71 | ops = [ |
| 72 | "add", |
| 73 | "divide", |
| 74 | "equal", |
| 75 | "greater", |
| 76 | "greater_equal", |
| 77 | "less", |
| 78 | "less_equal", |
| 79 | "logaddexp", |
| 80 | "maximum", |
| 81 | "minimum", |
| 82 | "multiply", |
| 83 | "power", |
| 84 | "subtract", |
| 85 | "logical_or", |
| 86 | "logical_and", |
| 87 | ] |
| 88 | for opname in ops: |
| 89 | with self.subTest(op=opname): |
| 90 | op = getattr(mx, opname) |
| 91 | x = mx.random.uniform(shape=(5,)) |
| 92 | y = mx.random.uniform(shape=(5,)) |
| 93 | out = mx.vmap(op)(x, y) |
| 94 | self.assertTrue(mx.array_equal(out, op(x, y))) |
| 95 | |
| 96 | x = mx.random.uniform(shape=(2, 4)) |
| 97 | y = mx.random.uniform(shape=(2, 4)) |
| 98 | out = mx.vmap(op)(x, y) |
| 99 | self.assertTrue(mx.array_equal(out, op(x, y))) |
| 100 | |
| 101 | out = mx.vmap(op, in_axes=(0, 0), out_axes=0)(x, y) |
| 102 | self.assertTrue(mx.array_equal(out, op(x, y))) |
| 103 | |
| 104 | y = mx.random.uniform(shape=(4, 2)) |
| 105 | out = mx.vmap(op, in_axes=(0, 1), out_axes=0)(x, y) |
| 106 | self.assertTrue(mx.array_equal(out, op(x, y.T))) |
| 107 | |
| 108 | out = mx.vmap(op, in_axes=(0, 1), out_axes=1)(x, y) |
| 109 | self.assertTrue(mx.array_equal(out, op(x, y.T).T)) |
| 110 | |
| 111 | def test_tree(self): |
| 112 | def my_fun(tree): |