rand converters

Author: apbose

py/torch_tensorrt/dynamo/conversion/ops_evaluators.py

@@ -58,10 +58,10 @@ def rand_validator(rand_node: Node) -> bool:
         )
         return False
     if layout is not None:
-        _LOGGER.debug(
-            f"Currently we don't support specifying layout, got {layout}."
-        )
-        return False    
+        _LOGGER.debug(f"Currently we don't support specifying layout, got {layout}.")
+        return False
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.rand.default)
 def aten_ops_rand(
     ctx: ConversionContext,
@@ -83,10 +83,10 @@ def randn_validator(randn_node: Node) -> bool:
         )
         return False
     if layout is not None:
-        _LOGGER.debug(
-            f"Currently we don't support specifying layout, got {layout}."
-        )
-        return False   
+        _LOGGER.debug(f"Currently we don't support specifying layout, got {layout}.")
+        return False
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.randn.default)
 def aten_ops_randn(
     ctx: ConversionContext,
@@ -108,10 +108,10 @@ def randperm_validator(randperm_node: Node) -> bool:
         )
         return False
     if layout is not None:
-        _LOGGER.debug(
-            f"Currently we don't support specifying layout, got {layout}."
-        )
-        return False   
+        _LOGGER.debug(f"Currently we don't support specifying layout, got {layout}.")
+        return False
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.randperm.default)
 def aten_ops_randperm(
     ctx: ConversionContext,
@@ -123,7 +123,5 @@ def aten_ops_randperm(
     device = kwargs.get("device", None)
     input = args[0]
     if not isinstance(input, int):
-        raise RuntimeError(
-            f"The input must be an integer"
-        )
-    return np.random.randperm(*args).to(device=device)
+        raise RuntimeError(f"The input must be an integer")
+    return np.random.randperm(*args).to(device=device)

tests/py/dynamo/conversion/test_rand_aten.py

@@ -14,12 +14,12 @@
     (
         "rand_two_dimension",
         (lambda shape: torch.ops.aten.rand(shape)),
-        [2,3],
+        [2, 3],
     ),
     (
         "rand_three_dimension",
         (lambda shape: torch.ops.aten.rand(shape)),
-        [2,3,4],
+        [2, 3, 4],
     ),
     (
         "randn_one_dimension",
@@ -29,12 +29,12 @@
     (
         "randn_two_dimension",
         (lambda shape: torch.ops.aten.randn(shape)),
-        [2,3],
+        [2, 3],
     ),
     (
         "randn_three_dimension",
         (lambda shape: torch.ops.aten.randn(shape)),
-        [2,3,4],
+        [2, 3, 4],
     ),
     (
         "randperm_one_case",
@@ -51,9 +51,9 @@
         (lambda x: torch.ops.aten.randperm(x)),
         [1500],
     ),
-
 ]
 
+
 class TestRandConverter(DispatchTestCase):
     @parameterized.expand(
         [
@@ -79,4 +79,4 @@ def forward(self, x):
 
 
 if __name__ == "__main__":
-    run_tests()
+    run_tests()