| 208 | self.assertTrue(mx.allclose(expected_grad, compiled_grad)) |
| 209 | |
| 210 | def test_vmap_compiled(self): |
| 211 | def simple_unary(x): |
| 212 | return -mx.exp(x) |
| 213 | |
| 214 | x = mx.array([[1.0, 2.0], [2.0, 3.0]]) |
| 215 | |
| 216 | expected_out = mx.vmap(simple_unary)(x) |
| 217 | out = mx.vmap(mx.compile(simple_unary))(x) |
| 218 | self.assertTrue(mx.allclose(expected_out, out)) |
| 219 | |
| 220 | def simple_binary(x, y): |
| 221 | return mx.abs(mx.exp(x + y) + y) |
| 222 | |
| 223 | x = mx.array([[1.0, -3.0], [0.5, -0.5]]) |
| 224 | y = mx.array([[2.0, -1.0], [0.25, -0.25]]) |
| 225 | |
| 226 | expected_out = mx.vmap(simple_binary)(x, y) |
| 227 | out = mx.vmap(mx.compile(simple_binary))(x, y) |
| 228 | self.assertTrue(mx.allclose(expected_out, out)) |
| 229 | |
| 230 | expected_out = mx.vmap(simple_binary, in_axes=(0, 1))(x, y) |
| 231 | out = mx.vmap(mx.compile(simple_binary), in_axes=(0, 1))(x, y) |
| 232 | self.assertTrue(mx.allclose(expected_out, out)) |
| 233 | |
| 234 | y = mx.array([0.25, -0.25]) |
| 235 | expected_out = mx.vmap(simple_binary, in_axes=(0, None))(x, y) |
| 236 | out = mx.vmap(mx.compile(simple_binary), in_axes=(0, None))(x, y) |
| 237 | self.assertTrue(mx.allclose(expected_out, out)) |
| 238 | |
| 239 | def simple_unary_outer(x): |
| 240 | x = mx.abs(x) |
| 241 | |
| 242 | @mx.compile |
| 243 | def simple_unary_inner(z): |
| 244 | return -mx.exp(x) |
| 245 | |
| 246 | return simple_unary_inner(x) |
| 247 | |
| 248 | expected_out = -mx.exp(mx.abs(x)) |
| 249 | out = mx.vmap(simple_unary_outer)(x) |
| 250 | self.assertTrue(mx.allclose(expected_out, out)) |
| 251 | |
| 252 | def test_vjp_vjp_compiled(self): |
| 253 | def simple_unary(x): |