.torch_interpreter.tracing

class experimental_experiment.torch_interpreter.tracing.CondCCOp

Cannot be imported from torch.ops.higher_order.cond (function cond overwrite submodule cond).

class experimental_experiment.torch_interpreter.tracing.CustomAttribute(root: CustomProxy, attr: str)

To trace attributes.

class experimental_experiment.torch_interpreter.tracing.CustomParameterProxy(tracer: TracerBase, node: Node, name, param)

A special proxy which lets “shape”, “size”, “dim”, and a few other attribute accesses pass through to the underlying module parameter object, so that conditional tests on these attributes will not throw exception during tracing.

class experimental_experiment.torch_interpreter.tracing.CustomProxy(node: Node, tracer: TracerBase | None = None)

Defines a custom proxy to trace the execution of a model and converts it into a fx graph. Works with CustomTracer.

classmethod cat(tensors: List[CustomProxy], dim: int = 0, *, out=None, axis: int | None = None) → CustomProxy

Implements cat for tensors.

instanceof(cls)

Tells if this proxy represents a specific class.

length()

Returns a proxy for the length.

class experimental_experiment.torch_interpreter.tracing.CustomProxyFloat(node: Node, tracer: TracerBase | None = None)

A proxy for a float.

instanceof(cls)

isinstance

class experimental_experiment.torch_interpreter.tracing.CustomProxyInt(node: Node, tracer: TracerBase | None = None)

A proxy for an integer.

instanceof(cls)

isinstance

class experimental_experiment.torch_interpreter.tracing.CustomTracer(autowrap_modules: ~typing.Tuple[ModuleType] = (<module 'math' (built-in)>, ), autowrap_functions: ~typing.Tuple[~typing.Callable, ...] = (), param_shapes_constant: bool = False)

Defines a custom tracer to trace the execution of a model and converts it into a fx graph. Works with CustomProxy.

::

from experimental_experiment.torch_interpreter.tracing import CustomTracer

graph = CustomTracer().trace(model)

create_arg(a: Any) → Argument

Overwrites this method to deal with more argument.

getattr(attr: str, attr_val: Any, parameter_proxy_cache: Dict[str, Any])

See torch.fx.Tracer.getattr().

classmethod graph_erase_node(graph: Graph, node: Node)

Removes a node and all predecessors with are only consumed by this one.

proxy(node: ~torch.fx.node.Node, cls: type[~experimental_experiment.torch_interpreter.tracing.CustomProxy] = <class 'experimental_experiment.torch_interpreter.tracing.CustomProxy'>) → Proxy

Overwrites this method to replace the default Proxy by CustomProxy.

register_callable(name: str, fn: Callable) → Node

Registers a function and return a unique name.

Parameters:

• name – prefix to prepend to the function name

• fn – function

Returns:

new_name

classmethod remove_inplace(graph: Graph, exported_program: ExportedProgram | None = None, verbose: int = 0, exc: bool = True) → int

Removes inplace operations.

Parameters:

• graph – graph to modify

• exported_program – if available, it is used in the error message to make it easier to trace the code source

• verbose – verbosity

• exc – raise an exception if not possible, other return -1

Returns:

number of inplace nodes removed

The most difficult pattern is the following:

%slice_11 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor]
    (args = (%clone, 0, 0, 9223372036854775807), kwargs = {})
%slice_12 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor]
    (args = (%slice_11, 1, 0, 9223372036854775807), kwargs = {})
%slice_13 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor]
    (args = (%slice_12, 2, 0, 9223372036854775807), kwargs = {})
%copy_ : [num_users=0] = call_function[target=torch.ops.aten.copy_.default]
    (args = (%slice_13, %masked_fill), kwargs = {})

classmethod remove_unnecessary_slices(graph: Graph) → int

Removes unnecessary slices:

Parameters:

graph – graph to modify

Returns:

number of inplace nodes removed

%slice_11 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor]
    (args = (%clone, 0, 0, 9223372036854775807), kwargs = {})

trace(root: Module | Callable[[...], Any], concrete_args: Dict[str, Any] | None = None, remove_inplace: bool = True, update_model_with_callable: bool = True) → Graph

Trace root and return the corresponding FX Graph representation. root can either be an nn.Module instance or a Python callable.

Note that after this call, self.root may be different from the root passed in here. For example, when a free function is passed to trace(), we will create an nn.Module instance to use as the root and add embedded constants to.

Parameters:

• root – Either a Module or a function to be traced through. Backwards-compatibility for this parameter is guaranteed.

• concrete_args – Concrete arguments that should not be treated as Proxies. This parameter is experimental and its backwards-compatibility is NOT guaranteed.

• remove_inplace – Removes inplace nodes

• update_model_with_attribute – in some cases (control flow), the model needs to be

Returns:

A Graph representing the semantics of the passed-in root.

experimental_experiment.torch_interpreter.tracing.replace_problematic_function_before_tracing()

Replaces function that cannot be traced with the default tracer such as torch.cat().

experimental_experiment.torch_interpreter.tracing.setitem_with_transformation(a, b, transformations)

Extended version of setitem to deal with inplace modification.