FX lowering

class torch_lattice.artifact.fx.LatticeArtifactInterpreter(module, builder)[source]

Bases: Interpreter

Lower an FX graph by interpreting it with symbolic lattice values.

Parameters:
call_module(target, args, kwargs)[source]

Execute a call_module node and return the result.

Parameters:

  • target (Callable[..., Any] | str) – The call target for this node. See Node for details on semantics

  • args (tuple[Any, ...]) – Tuple of positional args for this invocation

  • kwargs (dict[str, Any]) – Dict of keyword arguments for this invocation

Return type:

ArtifactValue

Return

Any: The value returned by the module invocation

Note

Backwards-compatibility for this API is guaranteed.

call_function(target, args, kwargs)[source]

Execute a call_function node and return the result.

Parameters:

  • target (Callable[..., Any] | str) – The call target for this node. See Node for details on semantics

  • args (tuple[Any, ...]) – Tuple of positional args for this invocation

  • kwargs (dict[str, Any]) – Dict of keyword arguments for this invocation

Return type:

Any

Return

Any: The value returned by the function invocation

Note

Backwards-compatibility for this API is guaranteed.

call_method(target, args, kwargs)[source]

Execute a call_method node and return the result.

Parameters:

  • target (Callable[..., Any] | str) – The call target for this node. See Node for details on semantics

  • args (tuple[Any, ...]) – Tuple of positional args for this invocation

  • kwargs (dict[str, Any]) – Dict of keyword arguments for this invocation

Return type:

Any

Return

Any: The value returned by the method invocation

Note

Backwards-compatibility for this API is guaranteed.

class torch_lattice.artifact.fx.LatticeTracer[source]

Bases: Tracer

FX tracer that preserves supported lattice modules and ops.

is_leaf_module(module, module_qualified_name)[source]

A method to specify whether a given nn.Module is a “leaf” module.

Leaf modules are the atomic units that appear in the IR, referenced by call_module calls. By default, Modules in the PyTorch standard library namespace (torch.nn) are leaf modules. All other modules are traced through and their constituent ops are recorded, unless specified otherwise via this parameter.

Parameters:

  • m (Module) – The module being queried about

  • module_qualified_name (str) – The path to root of this module. For example, if you have a module hierarchy where submodule foo contains submodule bar, which contains submodule baz, that module will appear with the qualified name foo.bar.baz here.

  • module (Module)

Return type:

bool

Note

Backwards-compatibility for this API is guaranteed.

torch_lattice.artifact.fx.lower_fx_artifact(builder, model, inputs=None)[source]
Return type:

TorchLatticeArtifactBuilder

Parameters: