Function create_program

extension/pybindings/test/make_test.py:151–186

Returns an executorch program based on ModuleAdd, along with inputs.

(
    eager_module: torch.nn.Module,
    et_config: Optional[ExecutorchBackendConfig] = None,
)

def create_program(
    eager_module: torch.nn.Module,
    et_config: Optional[ExecutorchBackendConfig] = None,
) -> Tuple[ExecutorchProgramManager, Tuple[Any, ...]]:
    """Returns an executorch program based on ModuleAdd, along with inputs."""

    # Trace the test module and create a serialized ExecuTorch program.
    # pyre-fixme[29]: `Union[torch._tensor.Tensor, torch.nn.modules.module.Module]`
    #  is not a function.
    inputs = eager_module.get_inputs()
    input_map = {}
    # pyre-fixme[29]: `Union[torch._tensor.Tensor, torch.nn.modules.module.Module]`
    #  is not a function.
    for method in eager_module.get_methods_to_export():
        input_map[method] = inputs

    class WrapperModule(torch.nn.Module):
        def __init__(self, fn):
            super().__init__()
            self.fn = fn

        def forward(self, *args, **kwargs):
            return self.fn(*args, **kwargs)

    exported_methods = {}
    # These cleanup passes are required to convert the `add` op to its out
    # variant, along with some other transformations.
    for method_name, method_input in input_map.items():
        wrapped_mod = WrapperModule(getattr(eager_module, method_name))
        exported_methods[method_name] = export(wrapped_mod, method_input, strict=True)

    exec_prog = to_edge(exported_methods).to_executorch(config=et_config)

    # Create the ExecuTorch program from the graph.
    exec_prog.dump_executorch_program(verbose=True)
    return (exec_prog, inputs)