pytorch - 💡(How to fix) Fix torch.export.save / torch.export.load does not support opaque type constants registered via register_opaque_type

from torch._library.opaque_object import register_opaque_type
from torch._opaque_base import OpaqueBase
class MyEngine(OpaqueBase):
    def __getstate__(self): ...
    def __setstate__(self, state): ...
register_opaque_type(MyEngine, typ="reference")

---

import torch
from torch._library.opaque_object import register_opaque_type
from torch._opaque_base import OpaqueBase
class MyEngine(OpaqueBase):
    def __init__(self, data):
        self.data = data
    def __getstate__(self):
        return {"data": self.data}
    def __setstate__(self, state):
        self.data = state["data"]
register_opaque_type(MyEngine, typ="reference")
@torch.library.custom_op("mylib::run_engine", mutates_args=())
def run_engine(x: torch.Tensor, engine: MyEngine) -> torch.Tensor:
    return x * engine.data
@run_engine.register_fake
def run_engine_fake(x: torch.Tensor, engine: MyEngine) -> torch.Tensor:
    return torch.empty_like(x)
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.engine = MyEngine(2.0)
    def forward(self, x):
        return torch.ops.mylib.run_engine(x, self.engine)
m = Model()
ep = torch.export.export(m, (torch.randn(4),))
print("Export succeeded")
print("Constants:", {k: type(v).__name__ for k, v in ep.constants.items()})
torch.export.save(ep, "/tmp/model.pt2")  # RuntimeError here
print("Save succeeded")

---

Export succeeded
Constants: {'lifted_custom_0': 'MyEngine'}
Traceback (most recent call last):
  File "repro.py", line 36, in <module>
    torch.export.save(ep, "/tmp/model.pt2")
  File ".../torch/export/__init__.py", line 274, in save
    package_pt2(
  File ".../torch/export/pt2_archive/_package.py", line 723, in package_pt2
    _package_exported_programs(
  File ".../torch/export/pt2_archive/_package.py", line 614, in _package_exported_programs
    constants_config = _package_constants(
                       ^^^^^^^^^^^^^^^^^^^
  File ".../torch/export/pt2_archive/_package.py", line 510, in _package_constants
    raise RuntimeError(f"Unsupported constant type: {type(constant)}")
RuntimeError: Unsupported constant type: <class '__main__.MyEngine'>

---

for constant_fqn, constant in exported_program.constants.items():
    if isinstance(constant, torch.Tensor):
        ...
    elif isinstance(constant, torch._C.ScriptObject):
        custom_objects.append((constant_fqn, constant))
    else:
        raise RuntimeError(f"Unsupported constant type: {type(constant)}")

---

elif path_name.startswith(CUSTOM_OBJ_FILENAME_PREFIX):
    constant_bytes = archive_reader.read_bytes(...)
    constants[constant_fqn] = torch._C._pickle_load_obj(constant_bytes)

---

elif attr_kind == _AttrKind.CONSTANT:
    if not isinstance(from_obj, (torch.Tensor, torch.ScriptObject)):
        raise AssertionError(
            f"expected torch.Tensor or torch.ScriptObject for CONSTANT attr_kind, got {type(from_obj)}"
        )