Summary:
Pull Request resolved: #186323
Removes # mypy: allow-untyped-defs from torch/__init__.py (and adds a matching pyrefly.toml sub-config) and reworks how the magic methods on SymInt, SymFloat, and SymBool are typed.
Previously each class carried a long hand-written list of placeholder stubs like def __add__(self, other) -> "SymInt": .... These didn't model the actual promotion rules — e.g. SymInt + float returns SymFloat, Tensor + SymT returns Tensor, and arithmetic on SymBool promotes to SymInt. Three new generic mixins encode those rules once and are reused across all three classes:
_SymTypingMagicAlsoBool[_PrimType, _BecomesIntPrimType, _BecomesIntSymType, _FloatPromotionType]— comparisons (==,!=,<,<=,>,>=) and the bool-promoting arithmetic ops (+,-,*, plus theirr-variants), with_FloatPromotionTypecontrolling whether the result can beSymFloat._SymTypingMagic[_PrimType, _SymType, _FloatPromotionType]— the rest of the numeric magic methods (abs,neg,floor,ceil,trunc,mod,lshift/rshift,pow_by_natural, the__sym_*__math wrappers,__int_truediv__/__int_floordiv__, etc.)._SymTypingMagicBitwise[_BitwiseLikeType]—__and__,__or__,__xor__and their reflected forms.
SymInt/SymFloat/SymBool then become small concrete subclasses parameterized over the right operand and result types (e.g. SymFloat's float-promotion type is _Never since it doesn't promote further).
Runtime side, sym_node.py gains a real xor entry in bitwise_ops, only_bool_magic_methods, and the dispatch table (with a new _sympy_xor), and also_bool_magic_methods is widened from just {"eq"} to the full comparison set {"eq", "ge", "gt", "le", "lt", "ne"} so the methods installed at runtime line up with what the new mixin advertises. A magic_methods_excl helper set is added for the remainder.
Other fallout from removing allow-untyped-defs:
PySymTypeis exported fromtorch._Cstubs and threaded into the asymmetric comparison-op signatures generated bytools/pyi/gen_pyi.py, soTensor.__lt__(SymInt)etc. type-check.SymInt.has_hint(),SymInt.hint,SymInt.constantmove from ad-hoc.node.*access to typedpropertyaccessors;definitely_true_hint/etc. are updated to use them.SymNode.shape_envisOptional, so callers insymbolic_shapes.pynow raise explicitAssertionError("shape_env should not be None")instead of relying on it being set. A couple ofmaybe_as_int()/maybe_as_float()call sites are tightened with walrus assignment to avoid calling the method twice.# pyrefly: ignore[missing-attribute]is added whereSymInt.nodeis duck-typed acrossSymNode | NestedIntNode | ConstantIntNode | LocalIntNodeand the attribute only exists on some.
Review order: start with torch/__init__.py — the three new mixins and the rewritten SymInt/SymFloat/SymBool definitions are where the design lives. Then torch/fx/experimental/sym_node.py for the runtime registration of xor and the broadened also_bool_magic_methods. Everything else is small call-site adjustments to satisfy the stricter typing.
Test Plan:
CI.
Authored with Claude.
Reviewed By: bobrenjc93
Differential Revision: D106389841
SOCIAL SHARE CARD GENERATOR