@@ -39,11 +39,70 @@ def _symbol_values(symbol: sympy.Symbol, shape_env: ShapeEnv) -> _ExactValues:
3939 return frozenset (sympy .Integer (value ) for value in range (lower , upper + 1 ))
4040
4141
42+ def _expr_symbols_to_values (
43+ expr : sympy .Basic ,
44+ shape_env : ShapeEnv ,
45+ ) -> dict [sympy .Symbol , _ExactValues ]:
46+ return {symbol : _symbol_values (symbol , shape_env ) for symbol in expr .free_symbols }
47+
48+
49+ def _try_expr_to_int (expr : sympy .Basic ) -> Optional [int ]:
50+ integer_value = _expr_to_int (expr )
51+ if integer_value is not None :
52+ return integer_value
53+
54+ try :
55+ return _expr_to_int (sympy .simplify (expr ))
56+ except (RecursionError , TypeError ):
57+ return None
58+
59+
60+ def _constant_expr_values (expr : sympy .Basic ) -> Optional [set [int ]]:
61+ if expr .free_symbols :
62+ return None
63+
64+ integer_value = _try_expr_to_int (expr )
65+ return {integer_value } if integer_value is not None else None
66+
67+
68+ def _evaluate_exact_values (
69+ expr : sympy .Basic ,
70+ shape_env : ShapeEnv ,
71+ ) -> _ExactValues :
72+ try :
73+ return sympy_interp (
74+ _ExactValueAnalysis ,
75+ _expr_symbols_to_values (expr , shape_env ),
76+ expr ,
77+ missing_handler = lambda symbol : _symbol_values (symbol , shape_env ),
78+ )
79+ except (RecursionError , TypeError ):
80+ return None
81+
82+
83+ def _exact_values_to_ints (exact_values : _ExactValues ) -> Optional [set [int ]]:
84+ if exact_values is None :
85+ return None
86+
87+ result : set [int ] = set ()
88+ for value in exact_values :
89+ integer_value = _try_expr_to_int (value )
90+ if integer_value is None :
91+ return None
92+ result .add (integer_value )
93+ return result
94+
95+
4296def _map_values (values : _ExactValues , fn ) -> _ExactValues :
4397 if values is None :
4498 return None
4599
46- result = {sympy .simplify (fn (value )) for value in values }
100+ result = set ()
101+ for value in values :
102+ try :
103+ result .add (fn (value ))
104+ except (RecursionError , TypeError ):
105+ return None
47106 if len (result ) > _MAX_SET_SIZE :
48107 return None
49108 return frozenset (result )
@@ -55,7 +114,13 @@ def _combine_values(lhs: _ExactValues, rhs: _ExactValues, fn) -> _ExactValues:
55114 if len (lhs ) * len (rhs ) > _MAX_SET_SIZE * _MAX_SET_SIZE :
56115 return None
57116
58- result = {sympy .simplify (fn (a , b )) for a in lhs for b in rhs }
117+ result = set ()
118+ for a in lhs :
119+ for b in rhs :
120+ try :
121+ result .add (fn (a , b ))
122+ except (RecursionError , TypeError ):
123+ return None
59124 if len (result ) > _MAX_SET_SIZE :
60125 return None
61126 return frozenset (result )
@@ -80,6 +145,12 @@ def mod(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues:
80145 return None
81146 return _combine_values (lhs , rhs , lambda a , b : sympy .Mod (a , b ))
82147
148+ @staticmethod
149+ def floordiv (lhs : _ExactValues , rhs : _ExactValues ) -> _ExactValues :
150+ if rhs is None or any (value == 0 for value in rhs ):
151+ return None
152+ return _combine_values (lhs , rhs , lambda a , b : sympy .floor (a / b ))
153+
83154 @staticmethod
84155 def pow (lhs : _ExactValues , rhs : _ExactValues ) -> _ExactValues :
85156 return _combine_values (lhs , rhs , lambda a , b : a ** b )
@@ -104,35 +175,15 @@ def evaluate_symbolic_expr_values(
104175) -> Optional [set [int ]]:
105176 """Return a best-effort finite set of possible integer values.
106177
107- The helper first relies on ``bound_sympy`` for cheap singleton detection.
108- When interval bounds are not precise enough , it falls back to a small
109- exact-set analysis over bounded symbols using ``sympy_interp``.
178+ The helper avoids ShapeEnv bound queries here because some exported dynamic
179+ expressions trigger very deep SymPy normalization. Instead , it relies on a
180+ small exact-set analysis over bounded symbols using ``sympy_interp``.
110181
111182 """
112- root_expr = sympy .simplify (
113- expr .node .expr if isinstance (expr , torch .SymInt ) else expr
114- )
115- value_range = shape_env .bound_sympy (root_expr )
116- if value_range .is_int and value_range .is_singleton ():
117- singleton = _expr_to_int (value_range .lower )
118- return {singleton } if singleton is not None else None
119-
120- exact_values = sympy_interp (
121- _ExactValueAnalysis ,
122- {
123- symbol : _symbol_values (symbol , shape_env )
124- for symbol in root_expr .free_symbols
125- },
126- root_expr ,
127- missing_handler = lambda symbol : _symbol_values (symbol , shape_env ),
128- )
129- if exact_values is None :
130- return None
183+ root_expr = expr .node .expr if isinstance (expr , torch .SymInt ) else expr
131184
132- result : set [int ] = set ()
133- for value in exact_values :
134- integer_value = _expr_to_int (sympy .simplify (value ))
135- if integer_value is None :
136- return None
137- result .add (integer_value )
138- return result
185+ constant_values = _constant_expr_values (root_expr )
186+ if constant_values is not None :
187+ return constant_values
188+
189+ return _exact_values_to_ints (_evaluate_exact_values (root_expr , shape_env ))
0 commit comments