SIGN IN SIGN UP

Handle Dynamo tracing of torch._C._infer_size (#185471)

Dynamo treats torch._C._infer_size as an in-graph callable. When the
arguments are torch.Size values, Dynamo represents the SizeVariable contents
as FX tuple-like aggregates, so fake execution and generated FX code end up
calling the Python C binding with tuples. That binding only accepts actual
torch.Size objects and fails with "expected a torch.Size as argument 1".
This breaks torch.compile on KL divergence for MultivariateNormal because its
KL helper broadcasts two batch shapes through torch._C._infer_size.

Add a Dynamo handler for torch._C._infer_size when both arguments are
SizeVariable. Fully constant sizes are evaluated through the eager C binding so
we preserve its exact result/error behavior. Symbolic sizes inline a small
polyfill mirroring ATen infer_size's right-aligned broadcast logic, including
its singleton handling and negative sentinel propagation. I avoided lowering to
torch.broadcast_shapes because that API rejects negative dimensions that
_infer_size can propagate when paired with singleton dimensions.

Fixes #155800
Generated by my agent

Benchmark Results:
Before this fix, the issue repro failed during Dynamo fake execution of
torch._C._infer_size, so there was no successful compiled timing. After this
fix, a 100-iteration warm benchmark of the issue KLModel on 32x10 inputs with
backend="eager" completed with:
- eager: mean 1.8966 ms, median 0.4318 ms, min 0.3199 ms
- compiled backend="eager": mean 1.0854 ms, median 0.4059 ms, min 0.3565 ms

Test Plan:
- python test/dynamo/test_functions.py FunctionTests.test_infer_size_with_size_inputs FunctionTests.test_infer_size_with_negative_one_dynamic_size
- python test/distributions/test_distributions.py TestKL.test_compile_kl_multivariate_normal
- issue KLModel repro with torch.compile(model, backend="eager")
- lintrunner -a

Pull Request resolved: https://github.com/pytorch/pytorch/pull/185471
Approved by: https://github.com/bobrenjc93
J
Jason Ansel committed
caca41c5881bd92dc249c983daf0812cadfd8d09
Parent: 0394392
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/23/2026, 10:05:26 PM