ciflow/trunk/186323: Type annotations (and runtime fallout) for torch/__init__.py (#186323)

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 their r-variants), with _FloatPromotionType controlling whether the result can be SymFloat.


• _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:

• PySymType is exported from torch._C stubs and threaded into the asymmetric comparison-op signatures generated by tools/pyi/gen_pyi.py, so Tensor.__lt__(SymInt) etc. type-check.


• SymInt.has_hint(), SymInt.hint, SymInt.constant move from ad-hoc .node.* access to typed property accessors; definitely_true_hint/etc. are updated to use them.


• SymNode.shape_env is Optional, so callers in symbolic_shapes.py now raise explicit AssertionError("shape_env should not be None") instead of relying on it being set. A couple of maybe_as_int() / maybe_as_float() call sites are tightened with walrus assignment to avoid calling the method twice.


• # pyrefly: ignore[missing-attribute] is added where SymInt.node is duck-typed across SymNode | NestedIntNode | ConstantIntNode | LocalIntNode and 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