Function while_loop

exir/control_flow.py:129–174

A higher order function returning the result based on executing body_fn until cond_fn returns False.

(
    cond_fn: Callable[..., torch.Tensor],
    body_fn: Callable[..., Tuple[torch.Tensor]],
    init_val: pytree.PyTree,
)

def while_loop(
    cond_fn: Callable[..., torch.Tensor],
    body_fn: Callable[..., Tuple[torch.Tensor]],
    init_val: pytree.PyTree,
) -> Union[Tuple[torch.Tensor], Value]:
    """
    A higher order function returning the result based on executing body_fn
    until cond_fn returns False.
    """
    flattened_inputs, _ = pytree.tree_flatten(init_val)
    if not all(isinstance(i, torch.Tensor) for i in flattened_inputs):
        raise ExportError(
            ExportErrorType.INVALID_INPUT_TYPE,
            f"control_flow.while_loop() expects all inputs values to be tensors, actual inputs: {init_val}",
        )

    with using_tracer(None):
        val = init_val
        while cond_fn(*val):
            val = body_fn(*val)

    flattened_outputs, _ = pytree.tree_flatten(val)
    if not all(isinstance(o, torch.Tensor) for o in flattened_outputs):
        raise ExportError(
            ExportErrorType.INVALID_OUTPUT_TYPE,
            f"control_flow.while_loop() expects all returned values to be tensors, actual outputs: {val}",
        )

    tracer = DispatchTracer.get()

    if tracer is None:
        return val

    gm_cond = _make_submodule(cond_fn, single_return=True)
    gm_body = _make_submodule(body_fn)

    proxies = tuple([unwrap_proxy(v) for v in flattened_inputs])

    proxy = tracer.create_proxy(
        "call_function",
        while_loop,
        (gm_cond, gm_body, proxies),
        {},
    )

    return tree_return(val, proxy, update_with_proxy)


def tracing_context(