rand converters

Author: apbose

py/torch_tensorrt/dynamo/conversion/ops_evaluators.py

@@ -47,3 +47,83 @@ def aten_ops_arange_start_step(
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     return np.arange(*args)
+
+
+def rand_validator(rand_node: Node) -> bool:
+    dtype = rand_node.kwargs.get("dtype", None)
+    layout = rand_node.kwargs.get("layout", None)
+    if dtype is not None:
+        _LOGGER.debug(
+            f"Currently we don't support specifying output dtype, got {dtype}."
+        )
+        return False
+    if layout is not None:
+        _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,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    device = kwargs.get("device", None)
+    return np.random.rand(*args).to(device=device)
+
+
+def randn_validator(randn_node: Node) -> bool:
+    dtype = randn_node.kwargs.get("dtype", None)
+    layout = randn_node.kwargs.get("layout", None)
+    if dtype is not None:
+        _LOGGER.debug(
+            f"Currently we don't support specifying output dtype, got {dtype}."
+        )
+        return False
+    if layout is not None:
+        _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,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    device = kwargs.get("device", None)
+    return np.random.randn(*args).to(device=device)
+
+
+def randperm_validator(randperm_node: Node) -> bool:
+    dtype = randperm_node.kwargs.get("dtype", None)
+    layout = randperm_node.kwargs.get("layout", None)
+    if dtype is not None:
+        _LOGGER.debug(
+            f"Currently we don't support specifying output dtype, got {dtype}."
+        )
+        return False
+    if layout is not None:
+        _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,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    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)

tests/py/dynamo/conversion/test_rand_aten.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+rand_ops = [
+    (
+        "rand_one_dimension",
+        (lambda shape: torch.ops.aten.rand(shape)),
+        [1],
+    ),
+    (
+        "rand_two_dimension",
+        (lambda shape: torch.ops.aten.rand(shape)),
+        [2,3],
+    ),
+    (
+        "rand_three_dimension",
+        (lambda shape: torch.ops.aten.rand(shape)),
+        [2,3,4],
+    ),
+    (
+        "randn_one_dimension",
+        (lambda shape: torch.ops.aten.randn(shape)),
+        [1],
+    ),
+    (
+        "randn_two_dimension",
+        (lambda shape: torch.ops.aten.randn(shape)),
+        [2,3],
+    ),
+    (
+        "randn_three_dimension",
+        (lambda shape: torch.ops.aten.randn(shape)),
+        [2,3,4],
+    ),
+    (
+        "randperm_one_case",
+        (lambda x: torch.ops.aten.randperm(x)),
+        [1],
+    ),
+    (
+        "randperm_two_case",
+        (lambda x: torch.ops.aten.randperm(x)),
+        [150],
+    ),
+    (
+        "randperm_three_case",
+        (lambda x: torch.ops.aten.randperm(x)),
+        [1500],
+    ),
+
+]
+
+class TestRandConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            (
+                rand_op[0],
+                rand_op[1],
+            )
+            for rand_op in rand_ops
+        ]
+    )
+    def test_rand(self, _, op, shape_or_input):
+        class TestModule(nn.Module):
+            def __init__(self, rand_op):
+                super().__init__()
+                self.rand_op = rand_op
+
+            def forward(self, x):
+                return self.rand_op(x)
+
+        inputs = [shape_or_input]
+        grid_model = TestModule(op)
+        self.run_test(grid_model, inputs)
+
+
+if __name__ == "__main__":
+    run_tests()