Setup LiveBook/Elixir to run Nx lib with Torchx (with GPU) on Macbook(Apple Silicon)

If you want to run Elixir Nx library for machine learning on Macbook/Mac Mini/Mac Studio with Metal (GPU) for Apple Silicon you can simple follow steps.

1. Setup and config Pytorch.

Install asdf if you don't have one.

Follow below script to install Python3 by asdf

asdf plugin-add python
asdf install python 3.x.x
asdf global python 3.x.x

pip3 install torch torchvision

Verify with simple Python script:

Run Python:

python

add text in Python:

import torch
x = torch.rand(5, 3)
print(x)

You can see output like:

tensor([[0.3380, 0.3845, 0.3217],
        [0.8337, 0.9050, 0.2650],
        [0.2979, 0.7141, 0.9069],
        [0.1449, 0.1132, 0.1375],
        [0.4675, 0.3947, 0.1426]])

Ctrl + D to exit.

2. Setup LiveBook in local for easy access local environment.

git clone https://github.com/livebook-dev/livebook.git
cd livebook
mix deps.get --only prod

# Run the Livebook server
MIX_ENV=prod mix phx.server
...
[Livebook] Application running at http://localhost:8080/?token=dqwkm3a6i3f7pwn6jfmug33322111

Access LiveBook from url in terminal.

LiveBook site

3. Create new LiveBook and config LiveBook.

Mix.install(
  [
    #...
    {:nx, "~> 0.7"},
    {:torchx, "~> 0.7"}
  ],
  config: [
    nx: [default_backend: {Torchx.Backend, device: :mps}]
  ]
)

Run setup for pull packages.

For case cannot load NIF (related with Torchx) module, try to run setup without cache in LiveBook.

Remember set compiler to default compiler for Axon if needed.

Add your script and have fun!