KEMBAR78
Add dlpack support for MPS device · Issue #153789 · pytorch/pytorch · GitHub
Skip to content

Add dlpack support for MPS device #153789

@stefanozampini

Description

@stefanozampini

🚀 The feature, motivation and pitch

It seems dlpack support for MPS tensors is missing.

$ cat test.py 
import torch

mps_device = torch.device('mps:0')
x = torch.ones(1, device=mps_device)
x.__dlpack_device__()
$ python test.py 
Traceback (most recent call last):
  File "/Users/zampins/Devel/petscml/test.py", line 5, in <module>
    x.__dlpack_device__()
  File "XXX/lib/python3.12/site-packages/torch/_tensor.py", line 1778, in __dlpack_device__
    raise ValueError(f"Unknown device type {torch_device_type} for Dlpack")
ValueError: Unknown device type mps for Dlpack
$ python -c 'import torch; print(torch.__version__, torch.version.git_version)'
2.7.0 134179474539648ba7dee1317959529fbd0e7f89

Alternatives

No response

Additional context

No response

cc @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen

Metadata

Metadata

Assignees

Labels

actionableenhancementNot as big of a feature, but technically not a bug. Should be easy to fixmodule: mpsRelated to Apple Metal Performance Shaders frameworktriagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions