[narendasan]

Summary

Models with complex values lost their exported signature or a complex output in three places. The review pointed out the first regression and the second limitation; this PR now fixes all three.

• Interior complex regions:
  • complex_decomposition_adapter retraces the graph with make_fx. That returns a plain CodeGen taking positional flat arguments, so the compiled module rejected its own keyword inputs. Llama 4's complex rotary embedding hits this:

    TypeError: forward() got an unexpected keyword argument 'input_ids'
    

  • The adapter now gives the graph a plain codegen while it retraces, so the retrace returns the graph's flat outputs. Afterwards it puts the exported program's pytree codegen back, both on the returned module and on the input module (which the legacy fallback reuses).
  • The first version of this PR restored the codegen and flattened the retraced output instead. That broke namedtuple and register_dataclass outputs: make_fx records building them as one node, so the restored codegen found one leaf instead of two.
• Complex inputs or outputs:
  • At the end of _insert_complex_io_adapters, the codegen was rebuilt from the graph's placeholder names (arg0_1, x_unpacked_complex). Such models could only be called by position.
  • It now keeps the exported program's codegen when lowering and partitioning carried it through (they do, with the change above). It rebuilds one only if it was lost.
• A complex output of a multi-output block:
  • The output boundary inserted view_as_complex only when the output was the block call itself. A block with several outputs returns each through a getitem.
  • So in a model returning {"prod": a * b, "energy": ...}, prod came back as a real [..., 2] tensor, on both complex paths. This predates the PR.
  • A getitem of a block now counts as that block's output.

Export → compile → call on both complex paths:

case	before this PR	first version	now
interior complex, tensor output	by position only on the decomposition path	✓	✓
interior complex, namedtuple or dataclass output	by position only on the decomposition path	✗	✓
complex input	by position only	by position only	✓
complex output	by position only on the decomposition path	same	✓
complex and real outputs of one block (dict)	prod returned real	same	✓

Tests

• New tests:
  • test_interior_complex_region_keeps_a_namedtuple_output (the adapter, by position and by keyword);
  • test_complex_input_complex_output_by_keyword[legacy|decomp];
  • test_complex_output_of_a_multi_output_block[legacy|decomp].
• With the earlier test_interior_complex_region_keeps_exported_signature:
  • all 6 pass;
  • on the first version of this PR, the 5 new ones fail;
  • before the PR, all 6 fail.
• pytest lowering/test_complex_decomposition_adapter.py lowering/test_complex_rewrite.py hlo/test_complex_ops.py hlo/test_complex_graph_break.py hlo/test_rope_embedding.py: 323 passed (the 318 of the first version, plus the 5 new ones).
• models/test_model_refit.py -k complex: 3 passed.
• trl-internal-testing/tiny-Llama4ForCausalLM runs and matches eager (cosine 1.0): one TensorRT engine, nothing left to PyTorch.
• Found by a Hugging Face architecture sweep. RTX 5090, TensorRT 11.3.

🤖 Generated with Claude Code

[shoumikhin]

Not blocking, for a follow-up: models with a complex input or output still cannot be called by keyword. At the end of _insert_complex_io_adapters, the codegen is rebuilt from the placeholder names (arg0_1, ...), which replaces the one this PR restores. Its comment says setting the codegen before partitioning crashes, which is no longer true with this change. Removing that rebuild fixed keyword calls for complex inputs and outputs in my runs, and the complex test files still passed.

[shoumikhin]

This breaks models whose output is a namedtuple or a dataclass registered with register_dataclass. make_fx records building that object as one call_function node, so output.args[0] is a single Node and tree_leaves returns one leaf. The restored out_spec expects two, so every call fails, also by position:

ValueError: treespec.unflatten(leaves): `leaves` has length 1 but the spec refers to a pytree that holds 2 items

On main the same model works when called by position. To reproduce, take the model from the new test and return NT(y, y * 2) with NT = collections.namedtuple("NT", ["a", "b"]).

An option that worked for me, without flattening: give gm a plain torch.fx.graph.CodeGen() only while decompose_complex_in_graph runs, so the retrace returns the flat graph outputs. Then put original_codegen back on gm and on the returned module. With that, namedtuple and dataclass outputs work again, and single, nested and dict outputs still work by position, by keyword and through save. A namedtuple case in the new test would keep this covered.

[shoumikhin]

Tested on A100: the new tests fail before this PR and pass with it, and the complex test files pass, also with #4896 merged in. Looks good to me. Two small things inline, neither blocks.

[shoumikhin]

This also wraps outputs that are already complex. Example: return torch.linalg.eig(x + 1) on the legacy path. eig runs in a PyTorch block, and the call now fails with "view_as_complex is not supported for complex tensors". It worked before this PR. A single-output eigvals already failed the same way before. The boundary should only convert outputs that were packed as real, not outputs that are already complex. That would fix both.

[shoumikhin]

This test still passes if I remove the plain CodeGen swap before the retrace. A namedtuple is also a tuple, so it gets rebuilt by accident. A dataclass output fails without the swap, so it would protect the swap better.