[export] serialize sympy.Exprs as ASTs instead of strings (#140084)

pianpwk · pytorchmergebot · commit d869344bc00b · 2024-11-20T07:44:25.000Z
Summary: The way we've been de/serializing sympy.Exprs is not roundtrippable in all cases (serialize by calling `str(expr)`, and deserialize by calling `sympy.sympify(expr_str)`). This has led to expressions being mathematically equivalent but structurally different, causing issues in ValueRanges. Example issue: #136797 This starts to deprecate the use of `expr_str` and stores expressions in AST format instead. For BC purposes, `expr_str` deserialization is still supported, but we will always serialize to `expr_ast`. We'll kill this once the serialization upgrader design is finalized and implemented. Test Plan: test_export Differential Revision: D65638757 Pull Request resolved: #140084 Approved by: https://github.com/angelayi
diff --git a/test/export/test_export.py b/test/export/test_export.py
@@ -8133,7 +8133,6 @@ def forward(self, x):
         export(f, (inputs,), dynamic_shapes=dynamic_shapes)
 
     @testing.expectedFailureRetraceabilityNonStrict
-    @testing.expectedFailureCppSerDes  # dynamic shape serialization
     def test_disable_forced_specializations_ok(self):
         # check that we don't force specialization, and defer to runtime asserts
         # with allow_complex_guards_as_runtime_asserts=True to successfully export
@@ -9192,9 +9191,7 @@ def forward(self, input1: torch.Tensor):
         inps = (torch.randn(1, 224, 768, device="cpu"),)
         export(Foo(), inps)
 
-    @testing.expectedFailureSerDer  # TODO(pianpwk): PowByNatural valuerange deserialization
-    @testing.expectedFailureCppSerDes  # TODO(pianpwk): PowByNatural valuerange deserialization
-    @testing.expectedFailureSerDerNonStrict
+    @testing.expectedFailureCppSerDes
     @testing.expectedFailureRetraceabilityNonStrict
     def test_dim_dynamic(self):
         dynamic = Dim.DYNAMIC
diff --git a/torch/_export/serde/schema.py b/torch/_export/serde/schema.py
@@ -8,7 +8,7 @@
 from torch._export.serde.union import _Union
 
 # NOTE: Please update this value if any modifications are made to the schema
-SCHEMA_VERSION = (8, 1)
+SCHEMA_VERSION = (8, 2)
 TREESPEC_VERSION = 1
 
 
@@ -62,13 +62,42 @@ class SymExprHint(_Union):
     as_bool: bool
 
 
+# A leaf node in a SymExprNode, containing a bool/float/int/sympy.Symbol.
+@dataclass(repr=False)
+class SymBase(_Union):
+    as_bool: bool
+    as_float: float
+    as_int: int
+    as_symbol: str
+
+
+# Represents an AST node in a sympy.Expr.
+# If not a leaf node, "target" is a string representing the operator,
+# and "args" is a list of child SymExprNodes.
+# If a leaf node, "target" is None, "args" is empty, and "base" is a SymBase
+# representing the leaf value.
+@dataclass(repr=False)
+class SymExprNode:
+    args: List["SymExprNode"] = field(default_factory=list)
+    target: Optional[str] = None
+    base: Optional[SymBase] = None
+
+
 # This is for storing the symbolic expressions behind symints/symfloats/symbools
-# For example, we can get something like
+# The deprecated "expr_str" field is easier to explain; we could store
 # SymExpr(expr_str="s0 + s1", hint=SymExprHint(as_int=4)
-# if we also have the hint that s0 and s1 are both 2.
+# for an expression where s0 == s1 == 2.
+# We're moving away from expr_str for roundtrippability, and now deserialize into
+# the "expr_ast" field, which is a tree representation of the expression,
+# containing a root SymExprNode.
+# While we're deprecating this, we'll store an empty string in "expr_str" for now,
+# and support deserialization for both "expr_str" and "expr_ast" fields.
+# We'll only serialize to "expr_ast".
+# TODO(pianpwk): implement upgrader & delete.
 @dataclass
 class SymExpr:
     expr_str: str
+    expr_ast: Optional[SymExprNode] = None
     hint: Optional[SymExprHint] = None
 
 
diff --git a/torch/_export/serde/schema.yaml b/torch/_export/serde/schema.yaml
@@ -1,5 +1,5 @@
 # @generated by update_schema.py
-# checksum<<8e27d48014d4ec1c773aef056c0c20b61bead54be8338e95d3347d3422472b9a>>
+# checksum<<74ae9f550efb42873fc58f07990cefe4da35d38c673beeacbe1e20808d9a6962>>
 Argument:
   kind: union
   fields:
@@ -346,6 +346,17 @@ SchemaVersion:
       type: int
     minor:
       type: int
+SymBase:
+  kind: union
+  fields:
+    as_bool:
+      type: bool
+    as_float:
+      type: float
+    as_int:
+      type: int
+    as_symbol:
+      type: str
 SymBool:
   kind: union
   fields:
@@ -365,6 +376,9 @@ SymExpr:
   fields:
     expr_str:
       type: str
+    expr_ast:
+      type: Optional[SymExprNode]
+      default: None
     hint:
       type: Optional[SymExprHint]
       default: None
@@ -377,6 +391,18 @@ SymExprHint:
       type: float
     as_bool:
       type: bool
+SymExprNode:
+  kind: struct
+  fields:
+    args:
+      type: List[SymExprNode]
+      default: '[]'
+    target:
+      type: Optional[str]
+      default: None
+    base:
+      type: Optional[SymBase]
+      default: None
 SymInt:
   kind: union
   fields:
@@ -437,5 +463,5 @@ UserOutputSpec:
       type: Argument
 SCHEMA_VERSION:
 - 8
-- 1
+- 2
 TREESPEC_VERSION: 1
diff --git a/torch/_export/serde/schema_check.py b/torch/_export/serde/schema_check.py
@@ -5,7 +5,7 @@
 import re
 import typing
 from enum import IntEnum
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, ForwardRef, List, Optional, Tuple, Union
 
 from torch._export.serde import schema
 from torch._export.serde.union import _Union
@@ -71,6 +71,8 @@ def dump_type(t) -> Tuple[str, str]:
                 )
             elif t == ():
                 return "()", ""
+            elif isinstance(t, ForwardRef):
+                return t.__forward_arg__, f"ForwardRef<{t.__forward_arg__}>"
             else:
                 raise AssertionError(f"Type {t} is not supported in export schema.")
 
diff --git a/torch/_export/serde/serialize.py b/torch/_export/serde/serialize.py
diff --git a/torch/csrc/utils/generated_serialization_types.h b/torch/csrc/utils/generated_serialization_types.h
