Flip deep learning (pytorch#2831)

Summary:
Pull Request resolved: pytorch#2831

X-link: facebookresearch/FBGEMM#26

Reviewed By: connernilsen

Differential Revision: D59653154

fbshipit-source-id: 2be7d6c6d540db34391bbd6e0dafe62eec185e82

Author: Maggie Moss

fbgemm_gpu/codegen/genscript/generate_backward_split.py

@@ -13,6 +13,14 @@
 from typing import List
 
 try:
+    # pyre-fixme[21]: Could not find name `ArgType` in
+    #  `deeplearning.fbgemm.fbgemm_gpu.codegen.genscript.optimizers`.
+    # pyre-fixme[21]: Could not find name `OptimItem` in
+    #  `deeplearning.fbgemm.fbgemm_gpu.codegen.genscript.optimizers`.
+    # pyre-fixme[21]: Could not find name `OptimizerArgsSet` in
+    #  `deeplearning.fbgemm.fbgemm_gpu.codegen.genscript.optimizers`.
+    # pyre-fixme[21]: Could not find name `generate_optimized_grad_sum_loop_access`
+    #  in `deeplearning.fbgemm.fbgemm_gpu.codegen.genscript.optimizers`.
     from .optimizers import *
     from .common import CodeTemplate
     from .optimizer_args import OptimizerArgsSet

fbgemm_gpu/codegen/genscript/optimizer_args.py

@@ -265,6 +265,7 @@ def make_function_schema_arg(ty: ArgType, name: str, default: Union[int, float])
         ArgType.PLACEHOLDER_TENSOR: tensor_arg,
         ArgType.INT: lambda x: int_arg(x, default=int(default)),
         ArgType.FLOAT: lambda x: float_arg(x, default=default),
+        # pyre-fixme[6]: For 2nd argument expected `int` but got `Union[float, int]`.
         ArgType.SYM_INT: lambda x: schema_sym_int_arg(x, default=default),
     }[ty](name)
 

fbgemm_gpu/test/quantize/fp8_rowwise_test.py

@@ -17,6 +17,9 @@
 from hypothesis import given, settings, Verbosity
 
 from . import common  # noqa E402
+
+# pyre-fixme[21]: Could not find name `open_source` in
+#  `deeplearning.fbgemm.fbgemm_gpu.test.quantize.common`.
 from .common import open_source
 
 if open_source:

fbgemm_gpu/test/quantize/fused_8bit_rowwise_test.py

@@ -15,6 +15,9 @@
 from hypothesis import assume, given, HealthCheck, settings
 
 from . import common  # noqa E402
+
+# pyre-fixme[21]: Could not find name `open_source` in
+#  `deeplearning.fbgemm.fbgemm_gpu.test.quantize.common`.
 from .common import (
     fused_rowwise_8bit_dequantize_2bytes_padding_scale_bias_first_reference,
     fused_rowwise_8bit_dequantize_reference,

fbgemm_gpu/test/quantize/fused_nbit_rowwise_test.py

@@ -15,6 +15,9 @@
 from hypothesis import assume, given, HealthCheck, settings
 
 from . import common  # noqa E402
+
+# pyre-fixme[21]: Could not find name `open_source` in
+#  `deeplearning.fbgemm.fbgemm_gpu.test.quantize.common`.
 from .common import (
     bytes_to_half_floats,
     fused_rowwise_nbit_quantize_dequantize_reference,

fbgemm_gpu/test/quantize/mixed_dim_int8_test.py

@@ -15,6 +15,9 @@
 from hypothesis import given, HealthCheck, settings
 
 from . import common  # noqa E402
+
+# pyre-fixme[21]: Could not find name `open_source` in
+#  `deeplearning.fbgemm.fbgemm_gpu.test.quantize.common`.
 from .common import open_source
 
 if open_source:

fbgemm_gpu/test/quantize/msfp_test.py

@@ -13,6 +13,9 @@
 from hypothesis import given, HealthCheck, settings
 
 from . import common  # noqa E402
+
+# pyre-fixme[21]: Could not find name `open_source` in
+#  `deeplearning.fbgemm.fbgemm_gpu.test.quantize.common`.
 from .common import open_source
 
 if open_source:

fbgemm_gpu/test/quantize/mx/common.py

@@ -142,6 +142,8 @@ def _get_format_params(  # noqa
 
     min_norm = _get_min_norm(ebits)
 
+    # pyre-fixme[6]: For 1st argument expected `ElemFormat` but got `Union[None,
+    #  ElemFormat, int]`.
     _FORMAT_CACHE[fmt] = (ebits, mbits, emax, max_norm, min_norm)
 
     return ebits, mbits, emax, max_norm, min_norm
@@ -308,6 +310,7 @@ def check_diff_quantize(
         raise IndexError
 
     # Convert to numpy
+    # pyre-fixme[9]: x has type `Tensor`; used as `Union[ndarray, Tensor]`.
     x = np.array(x) if type(x) is list else x
     x = x.cpu().numpy() if type(x) is torch.Tensor else x
     y1 = y1.detach().cpu().numpy()
@@ -510,11 +513,16 @@ def _quantize_elemwise_core(
         private_exp = None
 
     # Scale up so appropriate number of bits are in the integer portion of the number
+    # pyre-fixme[6]: For 3rd argument expected `Optional[int]` but got
+    #  `Optional[Tensor]`.
     out = _safe_lshift(out, bits - 2, private_exp)
 
+    # pyre-fixme[6]: For 3rd argument expected `RoundingMode` but got `str`.
     out = _round_mantissa(out, bits, round, clamp=False)
 
     # Undo scaling
+    # pyre-fixme[6]: For 3rd argument expected `Optional[int]` but got
+    #  `Optional[Tensor]`.
     out = _safe_rshift(out, bits - 2, private_exp)
 
     # Set values > max_norm to Inf if desired, else clamp them

fbgemm_gpu/test/quantize/mx4_test.py

@@ -21,6 +21,9 @@
 from hypothesis import given, settings, Verbosity
 
 from . import common  # noqa E402
+
+# pyre-fixme[21]: Could not find name `open_source` in
+#  `deeplearning.fbgemm.fbgemm_gpu.test.quantize.common`.
 from .common import open_source
 from .mx.common import (
     _get_format_params,
@@ -133,6 +136,8 @@ def fake_quantize_mx(
 
     # Undo tile reshaping
     if group_size:
+        # pyre-fixme[61]: `padded_shape` is undefined, or not always defined.
+        # pyre-fixme[61]: `orig_shape` is undefined, or not always defined.
         A = _undo_reshape_to_blocks(A, padded_shape, orig_shape, axes)
 
     return A

fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py

@@ -146,6 +146,7 @@ def generate_ssd_tbes(
         ssd_shards: int = 1,  # from SSDTableBatchedEmbeddingBags
         optimizer: OptimType = OptimType.EXACT_ROWWISE_ADAGRAD,
         cache_set_scale: float = 1.0,
+        # pyre-fixme[9]: pooling_mode has type `bool`; used as `PoolingMode`.
         pooling_mode: bool = PoolingMode.SUM,
         weights_precision: SparseType = SparseType.FP32,
         output_dtype: SparseType = SparseType.FP32,
@@ -239,6 +240,10 @@ def generate_ssd_tbes(
         if weights_precision == SparseType.FP16:
             emb_ref = [emb.float() for emb in emb_ref]
 
+        # pyre-fixme[7]: Expected `Tuple[SSDTableBatchedEmbeddingBags,
+        #  List[EmbeddingBag]]` but got `Tuple[SSDTableBatchedEmbeddingBags,
+        #  Union[List[Union[Embedding, EmbeddingBag]], List[Embedding],
+        #  List[EmbeddingBag]]]`.
         return emb, emb_ref
 
     def concat_ref_tensors(
@@ -724,7 +729,11 @@ def test_ssd_cache(
 
                 # pyre-fixme[16]: Optional type has no attribute `float`.
                 optim_state_r.add_(
-                    emb_r.weight.grad.float().to_dense().pow(2).mean(dim=1)
+                    # pyre-fixme[16]: `Optional` has no attribute `float`.
+                    emb_r.weight.grad.float()
+                    .to_dense()
+                    .pow(2)
+                    .mean(dim=1)
                 )
                 torch.testing.assert_close(
                     optim_state_t.float(),