After aten_cond creates an ONNX If node, the outputs were left without rank, type, or shape metadata in the graph builder — meaning _known_ranks, _known_types, and _known_shapes had no entries for those outputs. Any downstream optimization or graph-building pass that queries has_rank() / get_rank() on those results would get empty or wrong answers.

Changes

• experimental_experiment/torch_interpreter/_aten_functions.py — After g.make_node("If", ...) in aten_cond, propagate rank/type/shape from the true_graph local function builder to each If output. The builder already has complete metadata since both branches were fully interpreted before this call. Mirrors the same pattern used for submodule functions in interpreter.py.

if g.has_local_function(true_graph, g.local_domain, builder=True):
    fct_builder = g.get_local_function(true_graph, g.local_domain, builder=True)
    for bout, out in zip(fct_builder.output_names, outputs):
        if fct_builder.has_type(bout):
            g.set_type(out, fct_builder.get_type(bout))
        if fct_builder.has_device(bout):
            g.set_device(out, fct_builder.get_device(bout))
        if fct_builder.has_shape(bout):
            g.set_shape(out, fct_builder.get_shape(bout))
        elif fct_builder.has_rank(bout):
            g.set_rank(out, fct_builder.get_rank(bout))

Works correctly for single-output, multi-output, and dynamic-shape cases (builders store dynamic dimension strings such as 's77').

• _unittests/ut_torch_interpreter/test_onnx_export_control_flow.py — Added three targeted tests covering single-output cond, multi-output cond, and cond with dynamic shapes.

• Fixes fix missing rank after a node if (torch.cond) #573

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.