| 1303 | mx.eval(loss, grads) |
| 1304 | |
| 1305 | def test_compile_donates_input_buffer(self): |
| 1306 | mx.set_default_device(mx.cpu) |
| 1307 | |
| 1308 | def fun(x): |
| 1309 | return mx.sin(x) + 1 |
| 1310 | |
| 1311 | compiled_fn = mx.compile(fun) |
| 1312 | |
| 1313 | input = mx.arange(16, dtype=mx.float32) |
| 1314 | mx.eval(input) |
| 1315 | in_ptr = np.asarray(input, copy=False).__array_interface__["data"][0] |
| 1316 | |
| 1317 | out = compiled_fn(input) |
| 1318 | del input # Ensure the reference is dropped |
| 1319 | mx.eval(out) |
| 1320 | |
| 1321 | self.assertEqual( |
| 1322 | np.asarray(out, copy=False).__array_interface__["data"][0], in_ptr |
| 1323 | ) |
| 1324 | |
| 1325 | |
| 1326 | if __name__ == "__main__": |