InputGen: introduce argument generator

Reviewed By: SS-JIA

Differential Revision: D52047187

fbshipit-source-id: 48d0c65f78d6f3f2b1225e02fe4436211b2bf6ce

Author: manuelcandales

inputgen/argument/gen.py

@@ -0,0 +1,112 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Optional, Tuple
+
+import torch
+from inputgen.argument.engine import MetaArg
+from inputgen.variable.gen import VariableGenerator
+from inputgen.variable.space import VariableSpace
+from torch.testing._internal.common_dtype import floating_types, integral_types
+
+
+FLOAT_RESOLUTION = 8
+
+
+class TensorGenerator:
+    def __init__(
+        self, dtype: Optional[torch.dtype], structure: Tuple, space: VariableSpace
+    ):
+        self.dtype = dtype
+        self.structure = structure
+        self.space = space
+
+    def gen(self):
+        if self.dtype is None:
+            return None
+        vg = VariableGenerator(self.space)
+        min_val = vg.gen_min()
+        max_val = vg.gen_max()
+        if min_val == float("-inf"):
+            min_val = None
+        if max_val == float("inf"):
+            max_val = None
+        # TODO(mcandales): Implement a generator that actually supports any given space
+        return self.get_random_tensor(
+            size=self.structure, dtype=self.dtype, high=max_val, low=min_val
+        )
+
+    def get_random_tensor(self, size, dtype, high=None, low=None):
+        if low is None and high is None:
+            low = -100
+            high = 100
+        elif low is None:
+            low = high - 100
+        elif high is None:
+            high = low + 100
+        size = tuple(size)
+        if dtype == torch.bool:
+            if not self.space.contains(0):
+                return torch.full(size, True, dtype=dtype)
+            elif not self.space.contains(1):
+                return torch.full(size, False, dtype=dtype)
+            else:
+                return torch.randint(low=0, high=2, size=size, dtype=dtype)
+
+        if dtype in integral_types():
+            low = math.ceil(low)
+            high = math.floor(high) + 1
+        elif dtype in floating_types():
+            low = math.ceil(FLOAT_RESOLUTION * low)
+            high = math.floor(FLOAT_RESOLUTION * high) + 1
+        else:
+            raise ValueError(f"Unsupported Dtype: {dtype}")
+
+        if dtype == torch.uint8:
+            if not self.space.contains(0):
+                return torch.randint(low=max(1, low), high=high, size=size, dtype=dtype)
+            else:
+                return torch.randint(low=max(0, low), high=high, size=size, dtype=dtype)
+
+        t = torch.randint(low=low, high=high, size=size, dtype=dtype)
+        if not self.space.contains(0):
+            if high > 0:
+                pos = torch.randint(low=max(1, low), high=high, size=size, dtype=dtype)
+            else:
+                pos = torch.randint(low=low, high=0, size=size, dtype=dtype)
+            t = torch.where(t == 0, pos, t)
+
+        if dtype in integral_types():
+            return t
+        if dtype in floating_types():
+            return t / FLOAT_RESOLUTION
+
+
+class ArgumentGenerator:
+    def __init__(self, meta: MetaArg):
+        self.meta = meta
+
+    def gen(self):
+        if self.meta.optional:
+            return None
+        elif self.meta.argtype.is_tensor():
+            return TensorGenerator(
+                dtype=self.meta.dtype,
+                structure=self.meta.structure,
+                space=self.meta.value,
+            ).gen()
+        elif self.meta.argtype.is_tensor_list():
+            return [
+                TensorGenerator(
+                    dtype=self.meta.dtype[i],
+                    structure=self.meta.structure[i],
+                    space=self.meta.value,
+                ).gen()
+                for i in range(len(self.meta.dtype))
+            ]
+        else:
+            return self.meta.value

test/inputgen/test_argument_generator.py

@@ -0,0 +1,98 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from inputgen.argument.engine import MetaArg
+from inputgen.argument.gen import ArgumentGenerator, TensorGenerator
+from inputgen.argument.type import ArgType
+from inputgen.variable.solve import SolvableVariable
+
+
+class TestTensorGenerator(unittest.TestCase):
+    def test_gen(self):
+        v = SolvableVariable(float)
+        v.Ge(13)
+        v.Le(51)
+        tg = TensorGenerator(dtype=torch.float64, structure=(2, 3), space=v.space)
+        tensor = tg.gen()
+
+        self.assertEqual(tensor.shape, (2, 3))
+        self.assertEqual(tensor.dtype, torch.float64)
+        self.assertGreaterEqual(tensor.min(), 13)
+        self.assertLessEqual(tensor.max(), 51)
+
+    def test_zero_tensor(self):
+        v = SolvableVariable(float)
+        v.Eq(0)
+        tg = TensorGenerator(dtype=torch.float64, structure=(2, 3), space=v.space)
+        tensor = tg.gen()
+
+        self.assertEqual(tensor.shape, (2, 3))
+        self.assertEqual(tensor.dtype, torch.float64)
+        self.assertGreaterEqual(tensor.min(), 0)
+        self.assertLessEqual(tensor.max(), 0)
+
+
+class TestArgumentGenerator(unittest.TestCase):
+    def test_gen_optional(self):
+        m = MetaArg(argtype=ArgType.TensorOpt, optional=True)
+        tensor = ArgumentGenerator(m).gen()
+        self.assertEqual(tensor, None)
+
+    def test_gen_scalar(self):
+        m = MetaArg(argtype=ArgType.Scalar, value=True)
+        scalar = ArgumentGenerator(m).gen()
+        self.assertIs(scalar, True)
+
+    def test_gen_dim_list(self):
+        m = MetaArg(argtype=ArgType.DimList, structure=(2, 3))
+        dim_list = ArgumentGenerator(m).gen()
+        self.assertEqual(dim_list, [2, 3])
+
+    def test_gen_tensor(self):
+        v = SolvableVariable(float)
+        v.Ge(13)
+        v.Le(51)
+        m = MetaArg(
+            argtype=ArgType.Tensor, dtype=torch.float64, structure=(2, 3), value=v.space
+        )
+        tensor = ArgumentGenerator(m).gen()
+
+        self.assertEqual(tensor.shape, (2, 3))
+        self.assertEqual(tensor.dtype, torch.float64)
+        self.assertGreaterEqual(tensor.min(), 13)
+        self.assertLessEqual(tensor.max(), 51)
+
+    def test_gen_tensor_list(self):
+        v = SolvableVariable(float)
+        v.Ge(13)
+        v.Le(51)
+        m = MetaArg(
+            argtype=ArgType.TensorOptList,
+            dtype=[torch.float64, torch.int32, None],
+            structure=((2, 3), (3,), None),
+            value=v.space,
+        )
+        tensors = ArgumentGenerator(m).gen()
+
+        self.assertEqual(len(tensors), 3)
+        self.assertEqual(tensors[0].shape, (2, 3))
+        self.assertEqual(tensors[0].dtype, torch.float64)
+        self.assertGreaterEqual(tensors[0].min(), 13)
+        self.assertLessEqual(tensors[0].max(), 51)
+
+        self.assertEqual(tensors[1].shape, (3,))
+        self.assertEqual(tensors[1].dtype, torch.int32)
+        self.assertGreaterEqual(tensors[1].min(), 13)
+        self.assertLessEqual(tensors[1].max(), 51)
+
+        self.assertEqual(tensors[2], None)
+
+
+if __name__ == "__main__":
+    unittest.main()