MPS Sparse Support #129842
Description
🚀 The feature, motivation and pitch
I would like to be able to use torch.sparse on Metal, which currently breaks:
import torch
assert(torch.backends.mps.is_available())
mps_device = torch.device("mps")
crow_indices = torch.tensor([0, 2, 4], dtype=torch.int32).to(mps_device)
col_indices = torch.tensor([0, 1, 0, 1], dtype=torch.int32).to(mps_device)
values = torch.tensor([1, 2, 3, 4], dtype=torch.float32).to(mps_device)
csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, dtype=torch.float32)
with the following error:
% python3 file.py
Traceback (most recent call last):
File "file.py", line 9, in <module>
csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, dtype=torch.float32)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
NotImplementedError: Could not run 'new_compressed_tensor' from the 'mps:0' device.)
Alternatives
No response
Additional context
No response
cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer @jcaip @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen