ufmt reformatted and used numpy instead of math

Author: yxsamliu

fbgemm_gpu/bench/merge_embeddings_benchmark.py

@@ -16,7 +16,6 @@
 import numpy as np
 import tabulate
 import torch
-import math
 from fbgemm_gpu.split_embedding_configs import SparseType
 from fbgemm_gpu.split_table_batched_embeddings_ops_common import (
     BoundsCheckMode,
@@ -101,6 +100,7 @@ def generate_requests(
 
 # pyre-fixme[3]: Return type must be annotated.
 
+
 def _get_random_tensor(
     num_ads: int,
     embedding_dimension: int,
@@ -109,7 +109,7 @@ def _get_random_tensor(
     gpu_idx: int,
     include_quantization: bool,
     use_pitched: bool = True,
-    alignment: int = 256,   # alignment in bytes
+    alignment: int = 256,  # alignment in bytes
 ):
     device = torch.device(f"cuda:{gpu_idx}")
 
@@ -120,12 +120,14 @@ def _get_random_tensor(
 
         if use_pitched:
             width_bytes = width_elems * elem_size
-            pitch_bytes = math.ceil(width_bytes / alignment) * alignment
+            pitch_bytes = int(np.ceil(width_bytes / alignment) * alignment)
             pitch_elems = pitch_bytes // elem_size
             storage = torch.empty((num_ads, pitch_elems), dtype=dtype, device=device)
             result_tensor = storage[:, :width_elems]  # logical view
         else:
-            result_tensor = torch.randn(num_ads, width_elems, dtype=dtype, device=device)
+            result_tensor = torch.randn(
+                num_ads, width_elems, dtype=dtype, device=device
+            )
 
     elif data_type == "INT8":
         assert embedding_dimension % 2 == 0, "needs to align to 2 bytes for INT8"
@@ -135,12 +137,16 @@ def _get_random_tensor(
 
         if use_pitched:
             width_bytes = width_elems * elem_size
-            pitch_bytes = math.ceil(width_bytes / alignment) * alignment
+            pitch_bytes = int(np.ceil(width_bytes / alignment) * alignment)
             pitch_elems = pitch_bytes // elem_size
-            storage = torch.randint(0, 255, (num_ads, pitch_elems), dtype=dtype, device=device)
+            storage = torch.randint(
+                0, 255, (num_ads, pitch_elems), dtype=dtype, device=device
+            )
             result_tensor = storage[:, :width_elems]
         else:
-            result_tensor = torch.randint(0, 255, (num_ads, width_elems), dtype=dtype, device=device)
+            result_tensor = torch.randint(
+                0, 255, (num_ads, width_elems), dtype=dtype, device=device
+            )
 
     elif data_type == "INT4":
         assert embedding_dimension % 4 == 0, "needs to align to 2 bytes for INT4"
@@ -150,12 +156,16 @@ def _get_random_tensor(
 
         if use_pitched:
             width_bytes = width_elems * elem_size
-            pitch_bytes = math.ceil(width_bytes / alignment) * alignment
+            pitch_bytes = int(np.ceil(width_bytes / alignment) * alignment)
             pitch_elems = pitch_bytes // elem_size
-            storage = torch.randint(0, 255, (num_ads, pitch_elems), dtype=dtype, device=device)
+            storage = torch.randint(
+                0, 255, (num_ads, pitch_elems), dtype=dtype, device=device
+            )
             result_tensor = storage[:, :width_elems]
         else:
-            result_tensor = torch.randint(0, 255, (num_ads, width_elems), dtype=dtype, device=device)
+            result_tensor = torch.randint(
+                0, 255, (num_ads, width_elems), dtype=dtype, device=device
+            )
 
     else:
         raise ValueError

fbgemm_gpu/test/merge_pooled_embeddings_test.py

@@ -10,10 +10,11 @@
 
 import unittest
 from typing import Tuple
-import math
+
 import fbgemm_gpu
 
 import hypothesis.strategies as st
+import numpy as np
 import torch
 from hypothesis import given, settings, Verbosity
 
@@ -32,18 +33,23 @@
 
 typed_gpu_unavailable: tuple[bool, str] = gpu_unavailable
 
+
 def make_pitched_tensor(height, width, dtype, device, alignment=256):
-    elem_size = torch.finfo(dtype).bits // 8 if dtype.is_floating_point else torch.iinfo(dtype).bits // 8
+    elem_size = (
+        torch.finfo(dtype).bits // 8
+        if dtype.is_floating_point
+        else torch.iinfo(dtype).bits // 8
+    )
     width_bytes = width * elem_size
-    pitch_bytes = math.ceil(width_bytes / alignment) * alignment
+    pitch_bytes = int(np.ceil(width_bytes / alignment) * alignment)
     pitch_elems = pitch_bytes // elem_size
     storage = torch.randn((height, pitch_elems), dtype=dtype, device=device)
     view = storage[:, :width]  # logical shape
     return view.contiguous() if alignment == 0 else view  # return pitched view
 
 
-# @unittest.skipIf(*gpu_unavailable)
-# @unittest.skipIf(open_source, "Not supported in open source yet")
+@unittest.skipIf(*gpu_unavailable)
+@unittest.skipIf(open_source, "Not supported in open source yet")
 class MergePooledEmbeddingsTest(unittest.TestCase):
     # pyre-fixme[56]: Pyre was not able to infer the type of argument
     #  `hypothesis.strategies.integers($parameter$min_value = 1, $parameter$max_value =
@@ -140,11 +146,12 @@ def test_all_to_one_device(
             pitch = True
             if pitch:
                 inputs = [
-                make_pitched_tensor(10, 20, torch.float32, "cpu", alignment=256)
-                for _ in range(num_inputs)]
+                    make_pitched_tensor(10, 20, torch.float32, "cpu", alignment=256)
+                    for _ in range(num_inputs)
+                ]
             else:
                 inputs = [torch.randn(10, 20) for _ in range(num_inputs)]
-        
+
             cuda_inputs = [
                 input.to(f"cuda:{i % num_gpus}") for i, input in enumerate(inputs)
             ]