Is it possible to enable the GPU on an Apple M3 Max?

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.

Thanks for your answer. I’ve been using the recommendation from the documentation:
Mix.install([
{:exla, “~> 0.2”}
])

Nx.global_default_backend(EXLA.Backend)
But did not have any effect. I can give a try to Torchx, just in case. Thanks!

@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.

Ok, thank you @jonatanklosko! i’m digging into TorchX and mps. We’ll see if I have some luck