Function trace_call_delegate

exir/delegate.py:45–75

(proxy_mode, func_overload, lowered_module, *args)

    # pyre-ignore
    def trace_call_delegate(proxy_mode, func_overload, lowered_module, *args):
        # pyre-ignore
        def _unwrap_proxy(e):
            if not isinstance(e, (torch.Tensor, torch.SymInt, torch.SymFloat)):
                return e
            return get_proxy_slot(
                cast(torch.Tensor, e), proxy_mode.tracer, e, lambda e: e.proxy
            )

        if not is_lowered_module(lowered_module):
            raise ValueError(
                "executorch_call_delegate()'s first argument must be a LoweredBackendModule"
            )

        with disable_proxy_modes_tracing():
            out = call_delegate_cpu(lowered_module, *args)

        get_lowered_module_name(proxy_mode.tracer.root, lowered_module)

        node_args = (lowered_module, *args)
        proxy_args = pytree.tree_map(_unwrap_proxy, node_args)
        out_proxy = proxy_mode.tracer.create_proxy(
            "call_function",
            func_overload,
            proxy_args,
            {},
            name="executorch_call_delegate",
        )
        return track_tensor_tree(
            out, out_proxy, constant=None, tracer=proxy_mode.tracer
        )

    @executorch_call_delegate.py_impl(torch._C.DispatchKey.CompositeExplicitAutograd)
    # pyre-ignore