It uses [Nx](https://github.com/elixir-nx/nx/tree/main/nx#readme https://github.com/elixir-nx/nx/tree/main/nx#readme) (think of it as an equivalent to numpy/jax

It uses [Nx](https://github.com/elixir-nx/nx/tree/main/nx#readme https://github.com/elixir-nx/nx/tree/main/nx#readme) (think of it as an equivalent to numpy/jax) which uses Google XLA's as one of the possible backends to compile and run on the GPU.