Unable to use Nx indexed functions with Torchx :mps backend

[megaboy101]

I've been experimenting with Nx this past week and since I'm working on an M1 Mac I've been swapping between the :cpu and :mps devices in Torchx to compare relative performance.

I noticed when trying to use some of the indexed functions like Nx.indexed_add() and Nx.indexed_put() with the :mps device, I receive the following error:

Nx.indexed_put(Nx.tensor([0, 0, 0]), Nx.tensor([[1], [2]]), Nx.tensor([2, 4]))

** (RuntimeError) Torchx: All inputs of where should have same/compatible number of dims in NIF.where/3
    (torchx 0.5.1) lib/torchx.ex:445: Torchx.unwrap!/1
    (torchx 0.5.1) lib/torchx.ex:448: Torchx.unwrap_tensor!/2
    (torchx 0.5.1) lib/torchx/backend.ex:1403: Torchx.Backend.select/4
    (torchx 0.5.1) lib/torchx/backend.ex:536: Torchx.Backend.as_torchx_linear_indices/2
    (torchx 0.5.1) lib/torchx/backend.ex:504: Torchx.Backend.indexed/5

Note: This works fine then using the :cpu device, only the :mps device seems to cause these issues

[josevalim]

Yes, the :mps backend is experimental on both LibTorch itself and in Nx, so it is expected you will run into issues. A pull request is welcome!

[megaboy101]

Ah gotcha, ya I'd definitely be down to take a crack at a PR. Appreciate the clarification!

[josevalim]

Closing this, as it is a mps issue and support has to be added on LibTorch itself (which was most likely already done!).