Add torch.export.register_dataclass API (pytorch#109152)

`register_dataclass` allows dataclass to be used as valid input/output types of torch.export.export

Pull Request resolved: pytorch#109152
Approved by: https://github.com/ydwu4

docs/source/export.rst

@@ -554,6 +554,7 @@ API Reference
 .. autofunction:: constrain_as_value
 .. autofunction:: save
 .. autofunction:: load
+.. autofunction:: register_dataclass
 .. autoclass:: Constraint
 .. autoclass:: ExportedProgram
 

torch/export/__init__.py

@@ -18,6 +18,13 @@
 from torch.fx.passes.infra.pass_base import PassResult
 from torch.fx.passes.infra.pass_manager import PassManager
 
+from torch.utils._pytree import (
+    FlattenFunc,
+    FromDumpableContextFn,
+    ToDumpableContextFn,
+    UnflattenFunc,
+)
+
 
 __all__ = [
     "ArgumentKind",
@@ -33,6 +40,7 @@
     "dynamic_dim",
     "export",
     "load",
+    "register_dataclass",
     "save",
 ]
 
@@ -1006,6 +1014,7 @@ def specify_constraints(x):
     Acceptable types of inputs (for ``args`` and ``kwargs``) and outputs include:
 
     - Primitive types, i.e. ``torch.Tensor``, ``int``, ``float``, ``bool`` and ``str``.
+    - Dataclasses, but they must be registered by calling :func:`register_dataclass` first.
     - (Nested) Data structures comprising of ``dict``, ``list``, ``tuple``, ``namedtuple`` and
       ``OrderedDict`` containing all above types.
 
@@ -1129,3 +1138,37 @@ def load(
     return load(
         f, extra_files=extra_files, expected_opset_version=expected_opset_version
     )
+
+
+def register_dataclass(typ: Any) -> None:
+    """
+    Registers a dataclass as a valid input/output type for :func:`torch.export.export`.
+
+    Args:
+        typ: the dataclass type to register
+
+    Example::
+
+        @dataclass
+        class InputDataClass:
+            feature: torch.Tensor
+            bias: int
+
+        class OutputDataClass:
+            res: torch.Tensor
+
+        torch.export.register_dataclass(InputDataClass)
+        torch.export.register_dataclass(OutputDataClass)
+
+        def fn(o: InputDataClass) -> torch.Tensor:
+            res = res=o.feature + o.bias
+            return OutputDataClass(res=res)
+
+        ep = torch.export.export(fn, (InputDataClass(torch.ones(2, 2), 1), ))
+        print(ep)
+
+    """
+
+    from torch._export.utils import register_dataclass_as_pytree_node
+
+    return register_dataclass_as_pytree_node(typ)

torch/overrides.py

@@ -151,6 +151,7 @@ def get_ignored_functions() -> Set[Callable]:
         torch.export.dynamic_dim,
         torch.export.export,
         torch.export.load,
+        torch.export.register_dataclass,
         torch.export.save,
         torch.eye,
         torch.fft.fftfreq,