Add UInt4Tensor with tensor subclass#117557
Conversation
Summary: As a follow up from #117208, this PR added a UInt4Tensor in python, it can be used to construct a uint4 tensor and supports some basic operations like view, slice etc. We can extend this to support different quantized tensors as mentioned in the previous PR. Later * tensor factory support for uint4, and other sub byte dtypes * other sub byte tensor subclass support Test Plan: python test/test_tensors.py -k test_constructor Reviewers: Subscribers: Tasks: Tags: [ghstack-poisoned]
Summary: As a follow up from #117208, this PR added a UInt4Tensor in python, it can be used to construct a uint4 tensor and supports some basic operations like view, slice etc. We can extend this to support different quantized tensors as mentioned in the previous PR. Later * tensor factory support for uint4, and other sub byte dtypes * other sub byte tensor subclass support Test Plan: python test/test_tensors.py -k test_constructor Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 7461060 Pull Request resolved: #117557
| @@ -0,0 +1,5 @@ | |||
| from .uint1_7 import UInt4Tensor | |||
There was a problem hiding this comment.
Use torch._subclasses, it already exists
| @@ -0,0 +1,215 @@ | |||
| import torch | |||
There was a problem hiding this comment.
1_7 is a terrible name. How about sub_byte_uint or tiny_uint
There was a problem hiding this comment.
OK, I can use sub_byte_uint
ezyang
left a comment
There was a problem hiding this comment.
I'm not really sure what the point of putting this into pytorch/pytorch right now is? There is a lot of build out you need, and you are not making use of being in pytorch/pytorch in any essential way. Get further along the implementation before requesting review, IMO.
| [0x23, 0x45], | ||
| [0x23, 0x45], | ||
| ], dtype=torch.uint8)) | ||
| self.assertEqual(x[:, 2:6], expected) |
There was a problem hiding this comment.
Is this testing actually sufficient? Not convinced...
There was a problem hiding this comment.
you mean for the entire subclass or just the slicing operation? we do have a lot more tests that I'm planning to put in torchao repo
oh OK sure, I can build this in torchao first |
|
will add this later |
Stack from ghstack (oldest at bottom):
Summary:
As a follow up from #117208, this PR added a UInt4Tensor
in python, it can be used to construct a uint4 tensor and supports some basic operations like
view, slice etc.
We can extend this to support different quantized tensors as mentioned in the previous PR.
Later
Test Plan:
python test/test_tensors.py -k test_constructor
Reviewers:
Subscribers:
Tasks:
Tags: