pyre-fixmes for D75477355 - FBGEMM (#4390)

Summary:
Pull Request resolved: #4390

X-link: facebookresearch/FBGEMM#1461

Additional pyre-fixmes needed by pytorch change D75477355 that weren't caught by diff-time CI.

Reviewed By: bobrenjc93

Differential Revision: D77120313

fbshipit-source-id: 24f61082ce573c9db1973782efe2e4f94cfd44bf

Author: aorenste

fbgemm_gpu/fbgemm_gpu/uvm.py

@@ -21,6 +21,8 @@
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:cumem_utils")
 
 # Import all uvm enums from c++ library
+# pyre-fixme[6]: For 2nd argument expected `() -> List[Tuple[str, List[Tuple[str,
+#  int]]]]` but got `OpOverloadPacket`.
 create_enums(globals(), torch.ops.fbgemm.fbgemm_gpu_uvm_enum_query)
 
 

fbgemm_gpu/test/release/stable_release_test.py

@@ -144,7 +144,11 @@ def _test_stable_schema(self, version: str) -> None:
                     op_name = full_op_name.split(".")[3]
 
                     check_schema_compatibility_from_op_name(
-                        torch.ops.fbgemm, op_name, ref_schema_str
+                        # pyre-fixme[6]: For 1st argument expected `(...) -> Any`
+                        #  but got `_OpNamespace`.
+                        torch.ops.fbgemm,
+                        op_name,
+                        ref_schema_str,
                     )
 
     def test_backwards_compatibility(self) -> None:

fbgemm_gpu/test/sll/jagged_softmax_test.py

@@ -66,7 +66,9 @@ def test_triton_jagged_softmax(
             padded_x1 - (1.0 - presences.unsqueeze(2).to(padded_x1.dtype)) * 5e7
         )
         padded_ref = torch.nn.functional.softmax(
-            softmax_input.transpose(1, 2), dim=-1
+            # pyre-fixme[16]: `float` has no attribute `transpose`.
+            softmax_input.transpose(1, 2),
+            dim=-1,
         )  # [B, H, N]
         ref = torch.ops.fbgemm.dense_to_jagged(padded_ref.permute(0, 2, 1), [offsets])[
             0

fbgemm_gpu/test/tbe/eeg/tbe_indices_generator_test.py

@@ -58,5 +58,7 @@ def test_indices_generation(
 
         assert indices.shape == (num_indices,)
         assert indices.dtype == torch.int64
+        # pyre-fixme[6]: For 1st argument expected `Tensor` but got `bool`.
         assert not torch.any(indices > max_index)
+        # pyre-fixme[6]: For 1st argument expected `Tensor` but got `bool`.
         assert not torch.any(indices < 0)

fbgemm_gpu/test/tbe/ssd/ssd_utils_test.py

@@ -121,6 +121,8 @@ def test_masked_index_put(
             num_value_rows=num_indices,
             num_output_rows=num_output_rows,
             dtype=dtype,
+            # pyre-fixme[6]: For 7th argument expected `(Tensor, Tensor, Tensor,
+            #  Tensor, bool) -> Tensor` but got `OpOverloadPacket`.
             test_fn=torch.ops.fbgemm.masked_index_put,
             is_index_put=True,
             use_pipeline=use_pipeline,
@@ -154,6 +156,8 @@ def test_masked_index_select(
             num_value_rows=num_value_rows,
             num_output_rows=num_indices,
             dtype=dtype,
+            # pyre-fixme[6]: For 7th argument expected `(Tensor, Tensor, Tensor,
+            #  Tensor, bool) -> Tensor` but got `OpOverloadPacket`.
             test_fn=torch.ops.fbgemm.masked_index_select,
             is_index_put=False,
             use_pipeline=use_pipeline,