What's the right way to use torch dispatch?

[lichao0613]

If I want to enable our new device like "NPU" based on PyTorch, do I still need register a dispatch key(just like pytorch/pytorch#75863 does) when using torch_dispatch?

[albanD]

Hi,

Sorry for the long time to answer.
I think the question is best answered with an example so you can take a look at #36
Let me know what you think, how it works for you and if there are any shortcomings that you need help fixing!

[lichao0613]

Hi,

Sorry for the long time to answer. I think the question is best answered with an example so you can take a look at #36 Let me know what you think, how it works for you and if there are any shortcomings that you need help fixing!

Hi albanD.
Thanks a lot for the reply and demos.
en.... there are some questions when using the demos in #36.

How to rename Tensor.device.type as "my_device"？We can create a "my_device" tensor by wrapping a "meta" tensor, but we cannot modify its device_type. As we all know that "tensor.device.type == 'my_device' " is a common judgement statement for pytorch users.
the MyDeviceMode seems not work to me, such as t1 = torch.empty(4, device="my_device"). Is there anything I have missed?

[albanD]

How to rename Tensor.device.type as "my_device"？

When you call tensor.device that will call __torch_function__ for Tensor.device.__get__. You can then return anything you want from there.

the MyDeviceMode seems not work to me, such as t1 = torch.empty(4, device="my_device"). Is there anything I have missed?

Did you properly called enable_my_device() ?
If so, can you share a code sample to repro?

Note: leaving this issue open to track the conversation more easily on my end.