Add an example on how to add a new device out of tree #36
Conversation
new_device.py
Outdated
| @staticmethod | ||
| def __new__(cls, size, dtype, raw_data=None, requires_grad=False): | ||
| # Use a meta Tensor here to be used as the wrapper | ||
| return torch.Tensor._make_subclass(cls, torch.empty(size, dtype=dtype, device="meta"), require_grad=requires_grad) |
There was a problem hiding this comment.
No way to set the device alas!
There was a problem hiding this comment.
No, all the C++ code that sets "device" will make sure it is one of the ones in our hardcoded list :(
|
|
||
| def __repr__(self): | ||
| st = super().__repr__() | ||
| st = st.replace("device='meta'", "device='my_device'") |
There was a problem hiding this comment.
This looks more confusing than hell ha ha ha
There was a problem hiding this comment.
Yes, this most likely will need to be tweaked by real backends. This is just to be able to easily re-use the plain Tensor print.
new_device.py
Outdated
| except Exception as e: | ||
| print(e) | ||
| raise e | ||
| raise RuntimeError(f"No implementation for 'my_device' for {func}, {args}, {kwargs}") |
There was a problem hiding this comment.
If you use some of the tech from @eellison's fake tensor you should also be able to get the device property to report correctly as well
There was a problem hiding this comment.
I haven't looked into that one yet.
The problem is that even if that tries to create a torch.device based on the result from c++, that will fail.
|
cc @bdhirsh |
This is composed of:
Fixes #35