Skip to content

fix(dynamo): stop unpacking complex inputs before the refit output check - #4583

Open
shoumikhin wants to merge 2 commits into
pytorch:mainfrom
shoumikhin:fix/complex-refit-verify-inputs
Open

fix(dynamo): stop unpacking complex inputs before the refit output check#4583
shoumikhin wants to merge 2 commits into
pytorch:mainfrom
shoumikhin:fix/complex-refit-verify-inputs

Conversation

@shoumikhin

Copy link
Copy Markdown
Contributor

What is broken

Three tests are red on main:

  • test_complex_buffer_refit
  • test_dual_complex_buffer_refit
  • test_complex_buffer_with_real_param_refit
TypeError: view_as_real is only supported for complex tensors

Why

After refitting, refit_module_weights(verify_output=True) runs the lowered
PyTorch module and the compiled module on the same inputs and compares the
results. When an input is complex it first converted that input to a real tensor
of shape (..., 2), because the lowering pass used to rewrite a complex
placeholder into a real one.

That is no longer what the pass does. The placeholder stays complex and the
graph does the unpacking itself:

placeholder      arg0_1
call_function    aten.view_as_real.default   (arg0_1,)
call_function    aten.select.int             (view_as_real, 2, 0)
...

So the reference module got a real tensor and was then asked to unpack it a
second time, which is what raises.

Fix

Give both modules the inputs the caller passed. That is what the non-complex
path already did, so the special case goes away and the two paths become one.

Tested

On an H100 with the three tests above: they fail before this change and pass
after it. The whole file passes, 7 passed and 12 skipped, the skips being the
ones main already skips in this environment.

Three refit tests fail on main:

  test_complex_buffer_refit
  test_dual_complex_buffer_refit
  test_complex_buffer_with_real_param_refit

  TypeError: view_as_real is only supported for complex tensors

After refitting, `refit_module_weights(verify_output=True)` runs the lowered
PyTorch module and the compiled module on the same inputs and compares the
results. For a model with a complex input it first converted that input to a
real tensor of shape (..., 2), because the lowering pass used to rewrite the
complex placeholder to a real one.

It does not do that any more. The placeholder stays complex and the graph
starts with its own `view_as_real`, so the reference module was being handed a
real tensor and then asked to unpack it again, which is what raises.

Pass the caller's inputs to both modules, which is what the non-complex path
already did, so the two paths become one.
@meta-cla meta-cla Bot added the cla signed label Aug 26, 2026
@github-actions github-actions Bot added component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 26, 2026
@github-actions
github-actions Bot requested a review from cehongwang August 26, 2026 02:15
The repository lint job runs `black --check .` across the whole tree, so any
file that does not match the formatter fails CI for every open pull request,
not only the one that touched it. `tests/py/dynamo/conversion/test_cumsum_aten.py`
is currently not black-conformant on main, which turns the Python Linting check
red here.

Reformat that one file with black. This is a formatting-only change: two
statements that fit on a single line are un-wrapped. No test logic changes.

Verified by running `black --check .` on the full tree: all files pass.
@github-actions github-actions Bot added the component: tests Issues re: Tests label Aug 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant