suppress errors in pytorch

Differential Revision: D54963973

fbshipit-source-id: 0f3fb06bcc140c22ec99089d927a2641ced9aa37

Author: generatedunixname89002005307016

torchao/dtypes/nf4tensor.py

@@ -6,13 +6,17 @@
 import torch.nn.functional as F
 
 
+# pyre-fixme[5]: Global expression must be annotated.
 aten = torch.ops.aten
+# pyre-fixme[5]: Global expression must be annotated.
 c10d_functional = torch.ops.c10d_functional
 
 from typing import Any
+# pyre-fixme[5]: Global annotation cannot contain `Any`.
 NF4_OPS_TABLE: Dict[Any, Any] = {}
 
 
+# pyre-fixme[3]: Return type must be annotated.
 def same_metadata(a: "NF4Tensor", b: "NF4Tensor"):
     both_nf4 = isinstance(a, NF4Tensor) and isinstance(b, NF4Tensor)
     return (
@@ -22,9 +26,14 @@ def same_metadata(a: "NF4Tensor", b: "NF4Tensor"):
         and a.n_blocks == b.n_blocks
     )
 
+# pyre-fixme[3]: Return type must be annotated.
+# pyre-fixme[2]: Parameter must be annotated.
 def implements(aten_ops):
     """Use this decorator to implement a function for an aten op in __torch_dispatch__"""
 
+    # pyre-fixme[53]: Captured variable `aten_ops` is not annotated.
+    # pyre-fixme[3]: Return type must be annotated.
+    # pyre-fixme[2]: Parameter must be annotated.
     def decorator(func):
         for op in aten_ops:
             NF4_OPS_TABLE[op] = func
@@ -33,14 +42,20 @@ def decorator(func):
     return decorator
 
 @implements([torch.ops.aten.detach.default, torch.ops.aten.detach])
+# pyre-fixme[3]: Return type must be annotated.
+# pyre-fixme[2]: Parameter must be annotated.
 def noop_detach(func, *args, **kwargs):
     return args[0][0]
 
 @implements([torch.ops.aten._to_copy.default])
+# pyre-fixme[3]: Return type must be annotated.
+# pyre-fixme[2]: Parameter must be annotated.
 def _to_copy(func, *args, **kwargs):
     return args[0][0].get_original_weight().to(args[1]['dtype'])
 
 @implements([torch.ops.aten.to.dtype])
+# pyre-fixme[3]: Return type must be annotated.
+# pyre-fixme[2]: Parameter must be annotated.
 def to_dtype(func, *args, **kwargs):
     return args[0][0].get_original_weight().to(args[0][1])
 
@@ -50,11 +65,14 @@ def to_dtype(func, *args, **kwargs):
         aten.copy_.default,
     ]
 )
+# pyre-fixme[3]: Return type must be annotated.
+# pyre-fixme[2]: Parameter must be annotated.
 def copy_(func, *args, **kwargs):
     original: NF4Tensor = args[0][0]
     copy_in: torch.Tensor = args[0][1]
 
     # Base Case
+    # pyre-fixme[6]: For 2nd argument expected `NF4Tensor` but got `Tensor`.
     if same_metadata(original, copy_in):
         original_tensors = original.__tensor_flatten__()[0]
         for tensor_name in original_tensors:
@@ -76,6 +94,7 @@ def copy_(func, *args, **kwargs):
 @dataclass
 class SubclassTensorArgs:
     original_shape: torch.Size
+    # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
     original_strides: Tuple
     storage_offset: int
     dtype: torch.dtype
@@ -106,6 +125,7 @@ def get_block_absmax(inpt_tensor: torch.Tensor, block_size: int) -> torch.Tensor
 class NF4Tensor(torch.Tensor):
     """NF4Tensor class for converting a weight to the QLoRA NF4 format"""
 
+    # pyre-fixme[3]: Return type must be annotated.
     def __new__(
         cls,
         # Args related for base tensor construction
@@ -134,6 +154,7 @@ def __new__(
 
         """
 
+        # pyre-fixme[16]: `Tensor` has no attribute `_make_wrapper_subclass`.
         nf4tensor = torch.Tensor._make_wrapper_subclass(
             cls,
             tensor_meta.original_shape,
@@ -145,6 +166,7 @@ def __new__(
         )
         return nf4tensor
 
+    # pyre-fixme[3]: Return type must be annotated.
     def __init__(
         self,
         tensor_meta: SubclassTensorArgs,
@@ -169,6 +191,7 @@ def __init__(
 
     @classmethod
     @torch.no_grad()
+    # pyre-fixme[3]: Return type must be annotated.
     def from_tensor(
         cls,
         inpt_tensor: torch.Tensor,
@@ -281,6 +304,7 @@ def double_quantize_scalers(
             n_scaler_blocks, scaler_block_size
         )
 
+        # pyre-fixme[58]: `/` is not supported for operand types `int` and `Tensor`.
         quantization_factor = 256 / (2 * scaler_absmax)
         # Length equal to weight numel // block_size
         quantized_scaler_blocks = scaler_blocks * quantization_factor
@@ -290,6 +314,7 @@ def double_quantize_scalers(
         # This is needed to make sure that quantization_factor remains a repeated view of n_scaler_blocks
         # For some reason the 127/scaler_absmax realizes n_scaler entries when only n_scaler_blocks are needed
         # The following will grab the first entry for the n_scaler_blocks which is the same across the scaler_block_size
+        # pyre-fixme[16]: `float` has no attribute `__getitem__`.
         quantization_factor = quantization_factor[:, 0]
 
         return (
@@ -326,6 +351,7 @@ def dequantize_scalers(
 
     @staticmethod
     def convert_to_norm_float_weight(
+        # pyre-fixme[11]: Annotation `tensor` is not defined as a type.
         inpt_tensor: torch.Tensor, n_blocks: int, block_size: int, nf4: torch.tensor
     ) -> torch.Tensor:
         """Convert a tensor to the normalized float weight format"""
@@ -386,6 +412,7 @@ def get_original_weight(self) -> torch.Tensor:
 
     @staticmethod
     def quantize_tensor_nearest(
+        # pyre-fixme[11]: Annotation `float16` is not defined as a type.
         value: torch.float16, nf4: torch.Tensor
     ) -> torch.Tensor:
         """Quantize a float16 tensor to nf4 format to nearest and not rounded up"""
@@ -396,6 +423,10 @@ def quantize_tensor_nearest(
         return closest_nf4
 
     @staticmethod
+    # pyre-fixme[14]: `dequantize` overrides method defined in `TensorBase`
+    #  inconsistently.
+    # pyre-fixme[40]: Static method `dequantize` cannot override a non-static method
+    #  defined in `torch._C.TensorBase`.
     def dequantize(value: torch.Tensor, nf4: torch.Tensor) -> torch.Tensor:
         """Dequantize a nf4 value to float16 format"""
         # return nf4.index_select(0, value)
@@ -406,6 +437,8 @@ def unpack(
     ) -> Tuple[
         int, int, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Size
     ]:
+        # pyre-fixme[7]: Expected `Tuple[int, int, Tensor, Tensor, Tensor, Tensor,
+        #  Size]` but got `Tuple[int, int, int, Tensor, Tensor, Tensor, Tensor]`.
         return (
             self.block_size,
             self.n_blocks,
@@ -416,12 +449,15 @@ def unpack(
             self.quantized_data,
         )
 
+    # pyre-fixme[14]: `__repr__` overrides method defined in `Tensor` inconsistently.
+    # pyre-fixme[3]: Return type must be annotated.
     def __repr__(self):
         return f"Quantized Data: {self.quantized_data}\nScalers: {self.quantized_scalers}\n"
 
     def __str__(self):
         return f"NF4Tensor({self.shape}, {self.block_size})"
 
+    # pyre-fixme[3]: Return type must be annotated.
     def __tensor_flatten__(self):
         tensor_meta = SubclassTensorArgs(
             self.shape,
@@ -446,6 +482,10 @@ def __tensor_flatten__(self):
         ], ctx
 
     @staticmethod
+    # pyre-fixme[3]: Return type must be annotated.
+    # pyre-fixme[24]: Generic type `dict` expects 2 type parameters, use
+    #  `typing.Dict[<key type>, <value type>]` to avoid runtime subscripting errors.
+    # pyre-fixme[2]: Parameter must be annotated.
     def __tensor_unflatten__(inner_tensors: Dict, metadata, outer_size, outer_stride):
         assert len(inner_tensors) == 5, "Expected 5 inner tensors"
         return NF4Tensor(
@@ -461,17 +501,22 @@ def __tensor_unflatten__(inner_tensors: Dict, metadata, outer_size, outer_stride
         )
 
 
+    # pyre-fixme[3]: Return type must be annotated.
     def __str__(self):
         return self.to(torch.float32).__str__()
 
     @classmethod
+    # pyre-fixme[3]: Return type must be annotated.
+    # pyre-fixme[2]: Parameter must be annotated.
     def __torch_dispatch__(cls, func, types, args, kwargs=None):
         """TODO we are not supporting torch dispatch at the moment
         instead we have created a Autograd.Function to handle the linear
         """
         # All ops in the  NF4_OPS_TABLE expect NF4 Tensors as inputs
         # And don't support mixed tensor subclasses. This will trigger the handler for
         # the next type in the dispatch list
+        # pyre-fixme[3]: Return type must be annotated.
+        # pyre-fixme[2]: Parameter must be annotated.
         def allowed_subclasses(type):
             return (
                 issubclass(cls, type)
@@ -489,16 +534,25 @@ def allowed_subclasses(type):
         )
 
     # Do not force the Float8Tensor type on the returned tensor
+    # pyre-fixme[4]: Attribute must be annotated.
     __torch_function__ = torch._C._disabled_torch_function_impl
 
 class LinearNF4(torch.autograd.Function):
     @staticmethod
+    # pyre-fixme[14]: `forward` overrides method defined in `_SingleLevelFunction`
+    #  inconsistently.
+    # pyre-fixme[3]: Return type must be annotated.
+    # pyre-fixme[2]: Parameter must be annotated.
     def forward(ctx, input: torch.Tensor, weight: NF4Tensor):
         """Save the quantized nf4 weight for backward pass"""
         ctx.nf4_weight = weight
         return F.linear(input, weight.get_original_weight())
 
     @staticmethod
+    # pyre-fixme[14]: `backward` overrides method defined in `_SingleLevelFunction`
+    #  inconsistently.
+    # pyre-fixme[3]: Return type must be annotated.
+    # pyre-fixme[2]: Parameter must be annotated.
     def backward(ctx, grad_output):
         """The nf4 weight will never require grad so we can just return the grad_output @ weight.get_original_weight()"""
         weight: NF4Tensor = ctx.nf4_weight
@@ -514,6 +568,8 @@ def linear_nf4(input: torch.Tensor, weight: NF4Tensor) -> torch.Tensor:
     """
     return LinearNF4.apply(input, weight)
 
+# pyre-fixme[3]: Return type must be annotated.
+# pyre-fixme[2]: Parameter must be annotated.
 def to_nf4(tensor,
            block_size: int = 64,
            scaler_block_size: int = 256):