Add files using upload-large-folder tool

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/case.py

@@ -0,0 +1,188 @@
+import inspect
+import re
+import string
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from types import ModuleType
+
+import torch
+
+_TAGS: Dict[str, Dict[str, Any]] = {
+    "torch": {
+        "cond": {},
+        "dynamic-shape": {},
+        "escape-hatch": {},
+        "map": {},
+        "dynamic-value": {},
+        "operator": {},
+        "mutation": {},
+    },
+    "python": {
+        "assert": {},
+        "builtin": {},
+        "closure": {},
+        "context-manager": {},
+        "control-flow": {},
+        "data-structure": {},
+        "standard-library": {},
+        "object-model": {},
+    },
+}
+
+
+class SupportLevel(Enum):
+    """
+    Indicates at what stage the feature
+    used in the example is handled in export.
+    """
+
+    SUPPORTED = 1
+    NOT_SUPPORTED_YET = 0
+
+
+class ExportArgs:
+    __slots__ = ("args", "kwargs")
+
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+
+
+InputsType = Union[Tuple[Any, ...], ExportArgs]
+
+
+def check_inputs_type(x):
+    if not isinstance(x, (ExportArgs, tuple)):
+        raise ValueError(
+            f"Expecting inputs type to be either a tuple, or ExportArgs, got: {type(x)}"
+        )
+
+
+def _validate_tag(tag: str):
+    parts = tag.split(".")
+    t = _TAGS
+    for part in parts:
+        assert set(part) <= set(
+            string.ascii_lowercase + "-"
+        ), f"Tag contains invalid characters: {part}"
+        if part in t:
+            t = t[part]
+        else:
+            raise ValueError(f"Tag {tag} is not found in registered tags.")
+
+
+@dataclass(frozen=True)
+class ExportCase:
+    example_inputs: InputsType
+    description: str  # A description of the use case.
+    model: torch.nn.Module
+    name: str
+    extra_inputs: Optional[InputsType] = None  # For testing graph generalization.
+    # Tags associated with the use case. (e.g dynamic-shape, escape-hatch)
+    tags: Set[str] = field(default_factory=set)
+    support_level: SupportLevel = SupportLevel.SUPPORTED
+    dynamic_shapes: Optional[Dict[str, Any]] = None
+
+    def __post_init__(self):
+        check_inputs_type(self.example_inputs)
+        if self.extra_inputs is not None:
+            check_inputs_type(self.extra_inputs)
+
+        for tag in self.tags:
+            _validate_tag(tag)
+
+        if not isinstance(self.description, str) or len(self.description) == 0:
+            raise ValueError(f'Invalid description: "{self.description}"')
+
+
+_EXAMPLE_CASES: Dict[str, ExportCase] = {}
+_MODULES: Set[ModuleType] = set()
+_EXAMPLE_CONFLICT_CASES: Dict[str, List[ExportCase]] = {}
+_EXAMPLE_REWRITE_CASES: Dict[str, List[ExportCase]] = {}
+
+
+def register_db_case(case: ExportCase) -> None:
+    """
+    Registers a user provided ExportCase into example bank.
+    """
+    if case.name in _EXAMPLE_CASES:
+        if case.name not in _EXAMPLE_CONFLICT_CASES:
+            _EXAMPLE_CONFLICT_CASES[case.name] = [_EXAMPLE_CASES[case.name]]
+        _EXAMPLE_CONFLICT_CASES[case.name].append(case)
+        return
+
+    _EXAMPLE_CASES[case.name] = case
+
+
+def to_snake_case(name):
+    name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
+    return re.sub("([a-z0-9])([A-Z])", r"\1_\2", name).lower()
+
+
+def _make_export_case(m, name, configs):
+    if not issubclass(m, torch.nn.Module):
+        raise TypeError("Export case class should be a torch.nn.Module.")
+    m = m()
+
+    if "description" not in configs:
+        # Fallback to docstring if description is missing.
+        assert (
+            m.__doc__ is not None
+        ), f"Could not find description or docstring for export case: {m}"
+        configs = {**configs, "description": m.__doc__}
+    return ExportCase(**{**configs, "model": m, "name": name})
+
+
+def export_case(**kwargs):
+    """
+    Decorator for registering a user provided case into example bank.
+    """
+
+    def wrapper(m):
+        configs = kwargs
+        module = inspect.getmodule(m)
+        if module in _MODULES:
+            raise RuntimeError("export_case should only be used once per example file.")
+
+        assert module is not None
+        _MODULES.add(module)
+        normalized_name = to_snake_case(m.__name__)
+        module_name = module.__name__.split(".")[-1]
+        if module_name != normalized_name:
+            raise RuntimeError(
+                f'Module name "{module.__name__}" is inconsistent with exported program '
+                + f'name "{m.__name__}". Please rename the module to "{normalized_name}".'
+            )
+
+        case = _make_export_case(m, module_name, configs)
+        register_db_case(case)
+        return case
+
+    return wrapper
+
+
+def export_rewrite_case(**kwargs):
+    def wrapper(m):
+        configs = kwargs
+
+        parent = configs.pop("parent")
+        assert isinstance(parent, ExportCase)
+        key = parent.name
+        if key not in _EXAMPLE_REWRITE_CASES:
+            _EXAMPLE_REWRITE_CASES[key] = []
+
+        configs["example_inputs"] = parent.example_inputs
+        case = _make_export_case(m, to_snake_case(m.__name__), configs)
+        _EXAMPLE_REWRITE_CASES[key].append(case)
+        return case
+
+    return wrapper
+
+
+def normalize_inputs(x: InputsType) -> ExportArgs:
+    if isinstance(x, tuple):
+        return ExportArgs(*x)
+
+    assert isinstance(x, ExportArgs)
+    return x

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/autograd_function.py

@@ -0,0 +1,26 @@
+import torch
+
+from torch._export.db.case import export_case
+
+
+class MyAutogradFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        return x.clone()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output + 1
+
+
+@export_case(
+    example_inputs=(torch.randn(3, 2),),
+)
+class AutogradFunction(torch.nn.Module):
+    """
+    TorchDynamo does not keep track of backward() on autograd functions. We recommend to
+    use `allow_in_graph` to mitigate this problem.
+    """
+
+    def forward(self, x):
+        return MyAutogradFunction.apply(x)

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/cond_branch_class_method.py

@@ -0,0 +1,46 @@
+import torch
+
+from torch._export.db.case import export_case
+from functorch.experimental.control_flow import cond
+
+
+class MySubModule(torch.nn.Module):
+    def foo(self, x):
+        return x.cos()
+
+    def forward(self, x):
+        return self.foo(x)
+
+
+@export_case(
+    example_inputs=(torch.ones(3),),
+    tags={
+        "torch.cond",
+        "torch.dynamic-shape",
+    },
+)
+class CondBranchClassMethod(torch.nn.Module):
+    """
+    The branch functions (`true_fn` and `false_fn`) passed to cond() must follow these rules:
+      - both branches must take the same args, which must also match the branch args passed to cond.
+      - both branches must return a single tensor
+      - returned tensor must have the same tensor metadata, e.g. shape and dtype
+      - branch function can be free function, nested function, lambda, class methods
+      - branch function can not have closure variables
+      - no inplace mutations on inputs or global variables
+
+
+    This example demonstrates using class method in cond().
+
+    NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.subm = MySubModule()
+
+    def bar(self, x):
+        return x.sin()
+
+    def forward(self, x):
+        return cond(x.shape[0] <= 2, self.subm.forward, self.bar, [x])

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/cond_branch_nonlocal_variables.py

@@ -0,0 +1,63 @@
+import torch
+
+from torch._export.db.case import export_case
+from functorch.experimental.control_flow import cond
+
+
+@export_case(
+    example_inputs=(torch.ones(6),),
+    tags={
+        "torch.cond",
+        "torch.dynamic-shape",
+    },
+)
+class CondBranchNonlocalVariables(torch.nn.Module):
+    """
+    The branch functions (`true_fn` and `false_fn`) passed to cond() must follow these rules:
+    - both branches must take the same args, which must also match the branch args passed to cond.
+    - both branches must return a single tensor
+    - returned tensor must have the same tensor metadata, e.g. shape and dtype
+    - branch function can be free function, nested function, lambda, class methods
+    - branch function can not have closure variables
+    - no inplace mutations on inputs or global variables
+
+    This example demonstrates how to rewrite code to avoid capturing closure variables in branch functions.
+
+    The code below will not work because capturing closure variables is not supported.
+    ```
+    my_tensor_var = x + 100
+    my_primitive_var = 3.14
+
+    def true_fn(y):
+        nonlocal my_tensor_var, my_primitive_var
+        return y + my_tensor_var + my_primitive_var
+
+    def false_fn(y):
+        nonlocal my_tensor_var, my_primitive_var
+        return y - my_tensor_var - my_primitive_var
+
+    return cond(x.shape[0] > 5, true_fn, false_fn, [x])
+    ```
+
+    NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        my_tensor_var = x + 100
+        my_primitive_var = 3.14
+
+        def true_fn(x, y, z):
+            return x + y + z
+
+        def false_fn(x, y, z):
+            return x - y - z
+
+        return cond(
+            x.shape[0] > 5,
+            true_fn,
+            false_fn,
+            [x, my_tensor_var, torch.tensor(my_primitive_var)],
+        )

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/decorator.py

@@ -0,0 +1,26 @@
+import functools
+
+import torch
+
+from torch._export.db.case import export_case
+
+
+def test_decorator(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs) + 1
+
+    return wrapper
+
+
+@export_case(
+    example_inputs=(torch.ones(3, 2), torch.ones(3, 2)),
+)
+class Decorator(torch.nn.Module):
+    """
+    Decorators calls are inlined into the exported function during tracing.
+    """
+
+    @test_decorator
+    def forward(self, x, y):
+        return x + y

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/dynamic_shape_assert.py

@@ -0,0 +1,22 @@
+import torch
+
+from torch._export.db.case import export_case
+
+
+@export_case(
+    example_inputs=(torch.ones(3, 2),),
+    tags={"python.assert"},
+)
+class DynamicShapeAssert(torch.nn.Module):
+    """
+    A basic usage of python assertion.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        # assertion with error message
+        assert x.shape[0] > 2, f"{x.shape[0]} is greater than 2"
+        # assertion without error message
+        assert x.shape[0] > 1
+        return x

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/dynamic_shape_constructor.py

@@ -0,0 +1,19 @@
+import torch
+
+from torch._export.db.case import export_case
+
+
+@export_case(
+    example_inputs=(torch.ones(3, 2),),
+    tags={"torch.dynamic-shape"},
+)
+class DynamicShapeConstructor(torch.nn.Module):
+    """
+    Tensor constructors should be captured with dynamic shape inputs rather
+    than being baked in with static shape.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.ones(x.shape[0] * 2)

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/dynamic_shape_map.py

@@ -0,0 +1,23 @@
+import torch
+
+from torch._export.db.case import export_case
+from functorch.experimental.control_flow import map
+
+
+@export_case(
+    example_inputs=(torch.ones(3, 2), torch.ones(2)),
+    tags={"torch.dynamic-shape", "torch.map"},
+)
+class DynamicShapeMap(torch.nn.Module):
+    """
+    functorch map() maps a function over the first tensor dimension.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, xs, y):
+        def body(x, y):
+            return x + y
+
+        return map(body, xs, y)

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/dynamic_shape_round.py

@@ -0,0 +1,24 @@
+import torch
+
+from torch._export.db.case import export_case, SupportLevel
+from torch.export import Dim
+
+x = torch.ones(3, 2)
+dim0_x = Dim("dim0_x")
+
+@export_case(
+    example_inputs=(x,),
+    tags={"torch.dynamic-shape", "python.builtin"},
+    support_level=SupportLevel.NOT_SUPPORTED_YET,
+    dynamic_shapes={"x": {0: dim0_x}},
+)
+class DynamicShapeRound(torch.nn.Module):
+    """
+    Calling round on dynamic shapes is not supported.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x[: round(x.shape[0] / 2)]

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/fn_with_kwargs.py

@@ -0,0 +1,32 @@
+import torch
+
+from torch._export.db.case import export_case, ExportArgs, SupportLevel
+
+
+@export_case(
+    example_inputs=ExportArgs(
+        torch.randn(4),
+        (torch.randn(4), torch.randn(4)),
+        *[torch.randn(4), torch.randn(4)],
+        mykw0=torch.randn(4),
+        input0=torch.randn(4), input1=torch.randn(4)
+    ),
+    tags={"python.data-structure"},
+    support_level=SupportLevel.SUPPORTED,
+)
+class FnWithKwargs(torch.nn.Module):
+    """
+    Keyword arguments are not supported at the moment.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, pos0, tuple0, *myargs, mykw0, **mykwargs):
+        out = pos0
+        for arg in tuple0:
+            out = out * arg
+        for arg in myargs:
+            out = out * arg
+        out = out * mykw0
+        out = out * mykwargs["input0"] * mykwargs["input1"]
+        return out

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/nested_function.py

@@ -0,0 +1,27 @@
+import torch
+
+from torch._export.db.case import export_case
+
+
+@export_case(
+    example_inputs=(torch.ones(3, 2), torch.ones(2)),
+    tags={"python.closure"},
+)
+class NestedFunction(torch.nn.Module):
+    """
+    Nested functions are traced through. Side effects on global captures
+    are not supported though.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, a, b):
+        x = a + b
+        z = a - b
+
+        def closure(y):
+            nonlocal x
+            x += 1
+            return x * y + z
+
+        return closure(x)

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/null_context_manager.py

@@ -0,0 +1,26 @@
+import contextlib
+
+import torch
+
+from torch._export.db.case import export_case
+
+
+@export_case(
+    example_inputs=(torch.ones(3, 2),),
+    tags={"python.context-manager"},
+)
+class NullContextManager(torch.nn.Module):
+    """
+    Null context manager in Python will be traced out.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        """
+        Null context manager in Python will be traced out.
+        """
+        ctx = contextlib.nullcontext()
+        with ctx:
+            return x.sin() + x.cos()

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/optional_input.py

@@ -0,0 +1,19 @@
+import torch
+
+from torch._export.db.case import export_case, SupportLevel
+
+
+@export_case(
+    example_inputs=(torch.randn(2, 3),),
+    tags={"python.object-model"},
+    support_level=SupportLevel.NOT_SUPPORTED_YET,
+)
+class OptionalInput(torch.nn.Module):
+    """
+    Tracing through optional input is not supported yet
+    """
+
+    def forward(self, x, y=torch.ones(2, 3)):
+        if y is not None:
+            return x + y
+        return x

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/torch_sym_min.py

@@ -0,0 +1,17 @@
+import torch
+
+from torch._export.db.case import export_case, SupportLevel
+
+
+@export_case(
+    example_inputs=(torch.ones(3, 2),),
+    tags={"torch.operator"},
+    support_level=SupportLevel.NOT_SUPPORTED_YET,
+)
+class TorchSymMin(torch.nn.Module):
+    """
+    torch.sym_min operator is not supported in export.
+    """
+
+    def forward(self, x):
+        return x.sum() + torch.sym_min(x.size(0), 100)

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/examples/user_input_mutation.py

@@ -0,0 +1,18 @@
+import torch
+
+from torch._export.db.case import export_case, SupportLevel
+
+
+@export_case(
+    example_inputs=(torch.ones(3, 2),),
+    tags={"torch.mutation"},
+    support_level=SupportLevel.SUPPORTED,
+)
+class UserInputMutation(torch.nn.Module):
+    """
+    Directly mutate user input in forward
+    """
+
+    def forward(self, x):
+        x.mul_(2)
+        return x.cos()

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/db/logging.py

@@ -0,0 +1,2 @@
+def exportdb_error_message(case_name: str):
+    return ""

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/pass_infra/node_metadata.py

@@ -0,0 +1,32 @@
+from typing import Any, Dict, Set
+
+
+NodeMetadataValue = Any
+
+
+PROTECTED_KEYS: Set[str] = {
+    "val",
+    "stack_trace",
+    "nn_module_stack",
+    "debug_handle",
+    "tensor_meta",
+}
+
+
+class NodeMetadata:
+    def __init__(self, data: Dict[str, Any]) -> None:
+        self.data: Dict[str, Any] = data.copy()
+
+    def __getitem__(self, key: str) -> NodeMetadataValue:
+        return self.data[key]
+
+    def __setitem__(self, key: str, value: NodeMetadataValue) -> NodeMetadataValue:
+        if key in PROTECTED_KEYS:
+            raise RuntimeError(f"Could not override node key: {key}")
+        self.data[key] = value
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.data
+
+    def copy(self) -> "NodeMetadata":
+        return NodeMetadata(self.data.copy())

tuning-competition-baseline/.venv/lib/python3.11/site-packages/torch/_export/passes/__init__.py

@@ -0,0 +1 @@
+from .replace_view_ops_with_view_copy_ops_pass import ReplaceViewOpsWithViewCopyOpsPass