Function single_op_builder

backends/test/graph_builder.py:116–142

Create a graph module with a single op. Args: placeholders: Placeholders to be used as inputs to the GraphModule. op: The op to be inserted. args: The args to be passed to the op. kwargs: The kwargs to be passed to the op. Returns: A graph module wit

(
    placeholders: Sequence[Union[torch.Tensor, FakeTensor]],
    op: Target,
    args: Sequence[Argument],
    kwargs: Optional[dict[str, Argument]] = None,
)

def single_op_builder(
    placeholders: Sequence[Union[torch.Tensor, FakeTensor]],
    op: Target,
    args: Sequence[Argument],
    kwargs: Optional[dict[str, Argument]] = None,
) -> torch.fx.GraphModule:
    """Create a graph module with a single op.

    Args:
        placeholders: Placeholders to be used as inputs to the GraphModule.
        op: The op to be inserted.
        args: The args to be passed to the op.
        kwargs: The kwargs to be passed to the op.

    Returns:
        A graph module with a single op
    """
    builder = GraphBuilder()
    op_to_placeholder_dict = {
        p: builder.placeholder(f"p_{i}", p) for i, p in enumerate(placeholders)
    }
    proxy_args, proxy_kwargs = pytree.tree_map_only(
        (torch.Tensor, FakeTensor), lambda x: op_to_placeholder_dict[x], (args, kwargs)
    )
    node = builder.call_operator(op, proxy_args, proxy_kwargs)
    builder.output([node])
    return builder.get_graph_module()