Function test_cond_nested

tests/python/relax/test_frontend_from_exported_program.py:9253–9287

Nested cond: a cond inside one of the branches.

()

def test_cond_nested():
    """Nested cond: a cond inside one of the branches."""

    class CondNestedModel(Module):
        def forward(self, x):
            def true_fn(x):
                def inner_true(x):
                    return x * 2.0

                def inner_false(x):
                    return x * 3.0

                return torch.cond(x.sum() > 1, inner_true, inner_false, (x,))

            def false_fn(x):
                return x - 1.0

            return torch.cond(x.sum() > 0, true_fn, false_fn, (x,))

    example_args = (torch.randn(3, 4),)
    exported_program = export(CondNestedModel(), args=example_args)
    mod = from_exported_program(exported_program)

    assert isinstance(mod, tvm.IRModule)

    # Should have at least 4 branch functions (2 outer + 2 inner) plus main
    func_names = [gv.name_hint for gv in mod.get_global_vars()]
    branch_funcs = [n for n in func_names if n != "main"]
    assert len(branch_funcs) >= 4, (
        f"Expected at least 4 branch functions for nested cond, got {branch_funcs}"
    )
    # Verify no duplicate function names
    assert len(set(branch_funcs)) == len(branch_funcs), (
        f"Duplicate branch function names: {branch_funcs}"
    )


def test_affine_grid():