| 286 | self.assertTrue(mx.allclose(expected_out[0], out[0])) |
| 287 | |
| 288 | def test_transform_over_eval_compiled(self): |
| 289 | def outer(x): |
| 290 | y = mx.exp(mx.abs(x)) |
| 291 | mx.eval(y) |
| 292 | return y.sum() |
| 293 | |
| 294 | x = mx.array([2.0, -1.0, 0.5]) |
| 295 | dfdx = mx.grad(outer)(x) |
| 296 | |
| 297 | @mx.compile |
| 298 | def simple_unary(x): |
| 299 | return mx.exp(mx.abs(x)) |
| 300 | |
| 301 | def outer(x): |
| 302 | y = simple_unary(x) |
| 303 | mx.eval(y) |
| 304 | return y.sum() |
| 305 | |
| 306 | cdfdx = mx.grad(outer)(x) |
| 307 | self.assertTrue(mx.allclose(dfdx, cdfdx)) |
| 308 | |
| 309 | def test_compile_capture(self): |
| 310 | # Test update captured state outside compiled function |