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