INNER CODE UNIT · Python

Function

lean-dojo/TorchLean · torchlean/__init__.py:95

        class Function(torch.nn.Module):
            def forward(self, *inputs):
                return function(*inputs)

        model = Function()
    model = copy.deepcopy(model)
    examples = tuple(value.detach().clone().requires_grad_(value.requires_grad)
                     for value in examples)
    artifact = adapter.export_torchlean_graph_json(
        model, examples, str(directory / "graph.json"), require_torch_export=True)
    if execute:
        artifact["arguments"] = [adapter._constant(value) for value in examples]
    if compare:
        with torch.no_grad():
            outputs, _ = torch.utils._pytree.tree_flatten(model(*examples))
        if not all(isinstance(value, torch.Tensor) for value in outputs):
            raise ValueError("sample comparison requires tensor outputs")
        artifact["reference_outputs"] = [adapter._constant(value) for value in outputs]

View source record →