Note I’m not an expert so the following may not be completely accurate although I think the broad brushstrokes aren’t too far off.

'Nx` can operate with different backends and some of those backends, like tensorflow utilise GPUs.

To make that work efficiently, GPU code needs to be loaded onto the GPU and the data needs to be efficiently accessed on it too.

The way that happens is with the defn (define numerical) functions which very cleverly compile your Elixir code (a subset of the full language) into backend code. And if that backend executes on the GPU then you get the appropriate performance outcome. Much of Nx (and therefore libraries which depend on it like Axon) is written as defn.

Note no expert as well

elixir-nx can use different backends.. the torchx backend has beta support for the m1 gpu - and in a few weeks/months it will be super easy to get going - right now the install with gpu support is a bit cumbersome afaik..

https://twitter.com/meanderingstrm/status/1589570845537894402