Update on "Adding uint4 dtype implementation"

Summary:
This PR added some preliminary support for uint4 through tensor subclass and we'll continue to iterate on this

we plan to move the uint4 tensor subclass to core after it is more mature

Test Plan:
python test/dtypes/test_int4.py

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]

Author: jerryzh168

torchao/dtypes/uint4.py

@@ -82,7 +82,7 @@ def __new__(cls, elem, **kwargs):
         kwargs["requires_grad"] = False
         return torch.Tensor._make_wrapper_subclass(cls, up_size(elem.shape), dtype=torch.uint4, **kwargs)
 
-    def __init__(self, elem):
+    def __init__(self, elem, **kwargs):
         self.elem = elem
 
     @classmethod
@@ -242,12 +242,11 @@ def _dynamically_quantize_per_channel_int4(x, quant_min, quant_max, target_dtype
 
 class PerChannelSymmetricWeightUInt4Tensor(UInt4Tensor):
     @staticmethod
-    def __new__(cls, elem, scales):
-        return super().__new__(cls, elem)
+    def __new__(cls, elem, scales, **kwargs):
+        return super().__new__(cls, elem, **kwargs)
 
-    def __init__(self, elem, scales):
-        # super().__init__(elem)
-        self.elem = elem
+    def __init__(self, elem, scales, **kwargs):
+        super().__init__(elem, **kwargs)
         self.scales = scales