@hiramegl yeah, EXLA does not support Metal at the moment. The underlying project (XLA) does not have official support for Metal, there is a separate plugin from Apple, but it is closed-source and already outdated, so we wait for an official support. As for Torchx, the mps backend is somewhat experimental, so depending on what computation you run, some operations may not be supported, but worth a shot.

I tried this a while back and found It’s possible using Torchx and setting the device to mps.

However some APIs are not supported yet and resulted in errors.

For example, taking the sample at nx/exla/examples/regression.exs at main · elixir-nx/nx · GitHub at changing the default backend:

Nx.default_backend({Torchx.Backend, device: :mps})

give this error when running it:

** (ArgumentError) cannot perform operation across devices mps and cpu
    (torchx 0.6.4) lib/torchx.ex:467: anonymous fn/2 in Torchx.prepare_tensors_list!/2

I’d be interested to hear if anyone has managed to get it to work.