subclass_zoo icon indicating copy to clipboard operation
subclass_zoo copied to clipboard

What's the right way to use torch dispatch?

Open lichao0613 opened this issue 2 years ago • 3 comments

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

lichao0613 avatar May 16 '22 23:05 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 https://github.com/albanD/subclass_zoo/pull/36 Let me know what you think, how it works for you and if there are any shortcomings that you need help fixing!

albanD avatar Jun 02 '22 00:06 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!

Hi albanD. Thanks a lot for the reply and demos. en.... there are some questions when using the demos in https://github.com/albanD/subclass_zoo/pull/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?

lichao0613 avatar Jun 17 '22 08:06 lichao0613

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.

albanD avatar Jun 17 '22 14:06 albanD