experimental_experiment.torch_dynamo.partition

class experimental_experiment.torch_dynamo.partition.CustomOperatorSupport(unsupport_dict: Set[str] | None = None, verbose: int = 0)

is_node_supported(submodules: Mapping[str, Module], node: Node) → bool

Args:

submodules: mapping from module name to the module. This can be

retrieved by calling model.named_modules().

node: a Fx node that we want to determine whether it’s supported.

Returns:

is_supported: whether the arg node is supported.

experimental_experiment.torch_dynamo.partition.backend_partition_compile(graph_module: GraphModule, args: List[Tensor], support: OperatorSupport | None = None, allows_single_node_partition: bool = True, backend_function: Callable | None = None, use_aot_autograd: bool = True, decompositions=None, partition_fn=None, verbose: int = 1, dynamic: bool = False, full_graph: bool = True, **kwargs)

Partitions a graph module for any backend.