fix: moved uint2 to prototype folder

Author: melvinebenezer

test/dtypes/test_uint2.py

@@ -1,18 +1,9 @@
 from unittest import main
-
 import torch
 import torch.nn as nn
-
-from torch.testing._internal.common_quantization import (
-    QuantizationTestCase,
-)
-
-from torchao.dtypes.uint2 import (
-    BitnetTensor
-)
-from torchao.quantization.quant_api import (
-    _replace_with_custom_fn_if_matches_filter,
-)
+from torch.testing._internal.common_quantization import QuantizationTestCase
+from torchao.prototype.dtypes.uint2 import BitnetTensor
+from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter
 
 def _apply_weight_only_uint2_quant(model):
     def fn(mod):
@@ -25,7 +16,6 @@ def fn(mod):
         lambda mod, fqn: isinstance(mod, torch.nn.Linear),
     )
 
-
 class TestUInt2(QuantizationTestCase):
     def test_gpu_quant(self):
         device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -35,11 +25,6 @@ def test_gpu_quant(self):
             y_ref = m(x)
             _apply_weight_only_uint2_quant(m)
             y_wo = m(x)
-            # sqnr = compute_error(y_ref, y_wo)
-            # opt = torch.compile(m, fullgraph=True, mode="max-autotune")
-            # make sure it runs
-            # opt(x)
-
 
 if __name__ == "__main__":
     main()

torchao/dtypes/__init__.py

@@ -1,14 +1,12 @@
 from .nf4tensor import NF4Tensor, to_nf4
-from .uint2 import UInt2Tensor, BitnetTensor
+# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
 from .uint4 import UInt4Tensor
 from .aqt import AffineQuantizedTensor, to_aq
 from .float6_e3m2 import to_float6_e3m2, from_float6_e3m2
 
 __all__ = [
     "NF4Tensor",
     "to_nf4",
-    "UInt2Tensor",
-    "BitnetTensor",
     "UInt4Tensor"
     "AffineQuantizedTensor",
     "to_aq",

torchao/dtypes/uint2.py renamed to torchao/prototype/dtypes/uint2.py

@@ -2,7 +2,7 @@
 import torch._prims_common as utils
 import torch.utils._pytree as pytree
 from torch.library import impl, Library
-from .uint4 import qtensor_lib
+from ...dtypes.uint4 import qtensor_lib
 
 
 def down_size(size):
@@ -94,6 +94,7 @@ def pack_uint2(uint8_data: torch.Tensor) -> torch.Tensor:
         return output
 
 else:
+    # TODO: torch compile issue https://github.com/pytorch/pytorch/issues/127374 is fixed
     #@torch.compile
     def unpack_uint8_to_trinary2(uint8_data: torch.Tensor) -> torch.Tensor:
         # since we are using uint8 we will decode 4 entries per byte
@@ -150,7 +151,7 @@ def fill_defaults(args, n, defaults_tail):
     return r
 
 
-#qtensor_lib = Library("qtensors", "DEF")
+# qtensor_lib = Library("qtensors", "DEF")
 qtensor_lib.define(
     "quantize_per_tensor_uint2(Tensor input, float scale, int zero_point) -> Tensor"
 )