METAL Backend Support

#Question

Thanks for all your work on diffrax!

Has anyone gotten diffrax working with a METAL GPU backend? There may be a simple reason why this isn't possible that I'm not aware of.

from diffrax import diffeqsolve, Dopri5, ODETerm, SaveAt, PIDController

vector_field = lambda t, y, args: -y
term = ODETerm(vector_field)
solver = Dopri5()
saveat = SaveAt(ts=[0., 1., 2., 3.])
stepsize_controller = PIDController(rtol=1e-5, atol=1e-5)

diffeqsolve(term, solver, t0=0, t1=3, dt0=0.1, y0=1, saveat=saveat,
                  stepsize_controller=stepsize_controller)

Returns:

ValueError: `EmitPythonCallback` not supported on METAL backend.
Run `reticulate::py_last_error()` for details.
5.
stop(<environment>)
4.
pure_callback_lowering at callback.py#141
3.
_call at _jit.py#98
2.
__call__ at _jit.py#102
1.
diffrax$diffeqsolve(term, solver, t0 = 0, t1 = 1, dt0 = 0.1,
y0 = y0, stepsize_controller = stepsize_controller, throw = F)

Start by reproducing the reported failure through diffrax.diffeqsolve with the METAL backend and the example configuration in the issue. Investigate the EmitPythonCallback limitation and determine whether METAL support is feasible; done means the example runs successfully without that error, with any required scope documented.