Device type for torch MPS tensor

[awni]

Should there be a supported device type for a mps tensor in torch to convert to a nb::ndarray?

I tried the metal device type but it did not match.

E.g. trying to bind a function which can accept torch.tensor([1,2,3]).to("mps") that as an nb::ndarray.

[wjakob]

If nb::device::metal does not work, what does PyTorch specify?

[awni]

what does PyTorch specify?

Meaning what's the device in like the dlpack tensor?

[wjakob]

Exactly. I would put some printfs into the code in nb_ndarray.cpp that accepts foreign tensors/ndarrays to see where it goes wrong. See ndarray_import.

[awni]

Well maybe this is the issue? Looks like PyTorch doesn't support mps with dlpack.

import torch
t = torch.tensor([1,2,3]).to("mps")
t.__dlpack__()

RuntimeError: Cannot pack tensors on mps:0

[wjakob]

Probably just a missing case in a switch statement. I would file an issue and maybe PR there, the team has been willing to merge DLPack-related improvements in the past.

[awni]

Actually looks like they fixed it recently switching from PyTorch 2.8.0 to 2.9.0 works.