dmx.compressor.fx.transformer.nodedict_transformer

Classes

NodeDictTransformer(module)

A transformer that creates a dict contaning mapping between target of node and the node itself

class dmx.compressor.fx.transformer.nodedict_transformer.NodeDictTransformer(module: GraphModule)

Bases: Transformer

A transformer that creates a dict contaning mapping between target of node and the node itself

Parameters:: module (fx.GraphModule) – the module to create the mapping

module

the module to create the mapping

Type:: fx.GraphModule

nodeDict

dictionary to store the mapping between occurance order of the node and the node

Type:: dict

call_function(target: Target, args: Tuple[tuple[Argument, ...] | Sequence[Argument] | Mapping[str, Argument] | slice | range | Node | str | int | float | bool | complex | dtype | Tensor | device | memory_format | layout | OpOverload | SymInt | SymBool | SymFloat | None, ...], kwargs: Dict[str, Any]) → Any

Execute a call_function node. Adds mapping between target of the node and the node itself to self.nodeDict.

Parameters:

target (Target) – The call target for this node. See Node for details on semantics
args (Tuple) – Tuple of positional args for this invocation
kwargs (Dict) – Dict of keyword arguments for this invocation

Returns:

A Proxy containing the node and the tracer of the new graph

call_method(target: Target, args: Tuple[tuple[Argument, ...] | Sequence[Argument] | Mapping[str, Argument] | slice | range | Node | str | int | float | bool | complex | dtype | Tensor | device | memory_format | layout | OpOverload | SymInt | SymBool | SymFloat | None, ...], kwargs: Dict[str, Any]) → Any

Execute a call_method node. Adds mapping between target of the node and the node itself to self.nodeDict.

Parameters:

target (Target) – The call target for this node. See Node for details on semantics
args (Tuple) – Tuple of positional args for this invocation
kwargs (Dict) – Dict of keyword arguments for this invocation

Returns:

A Proxy containing the node and the tracer of the new graph

call_module(target: Target, args: Tuple[tuple[Argument, ...] | Sequence[Argument] | Mapping[str, Argument] | slice | range | Node | str | int | float | bool | complex | dtype | Tensor | device | memory_format | layout | OpOverload | SymInt | SymBool | SymFloat | None, ...], kwargs: Dict[str, Any]) → Any

Execute a call_module node. Adds mapping between target of the node and the node itself to self.nodeDict.

Parameters:

target (Target) – The call target for this node. See Node for details on semantics
args (Tuple) – Tuple of positional args for this invocation
kwargs (Dict) – Dict of keyword arguments for this invocation

Returns:

A Proxy containing the node and the tracer of the new graph

get_attr(target: Target, args: Tuple[tuple[Argument, ...] | Sequence[Argument] | Mapping[str, Argument] | slice | range | Node | str | int | float | bool | complex | dtype | Tensor | device | memory_format | layout | OpOverload | SymInt | SymBool | SymFloat | None, ...], kwargs: Dict[str, Any]) → Proxy

Execute a get_attr node. Adds mapping between target of the node and the node itself to self.nodeDict.

Parameters:

target (Target) – The call target for this node. See Node for details on semantics
args (Tuple) – Tuple of positional args for this invocation
kwargs (Dict) – Dict of keyword arguments for this invocation

Returns:

A Proxy containing the node and the tracer of the new graph

output(target: Target, args: Tuple[tuple[Argument, ...] | Sequence[Argument] | Mapping[str, Argument] | slice | range | Node | str | int | float | bool | complex | dtype | Tensor | device | memory_format | layout | OpOverload | SymInt | SymBool | SymFloat | None, ...], kwargs: Dict[str, Any]) → Any

Execute a ouptput node. Adds mapping between target of the node and the node itself to self.nodeDict.

Parameters:

target (Target) – The call target for this node. See Node for details on semantics
args (Tuple) – Tuple of positional args for this invocation
kwargs (Dict) – Dict of keyword arguments for this invocation

Returns:

A Proxy containing the node and the tracer of the new graph

placeholder(target: Target, args: Tuple[tuple[Argument, ...] | Sequence[Argument] | Mapping[str, Argument] | slice | range | Node | str | int | float | bool | complex | dtype | Tensor | device | memory_format | layout | OpOverload | SymInt | SymBool | SymFloat | None, ...], kwargs: Dict[str, Any]) → Proxy

Execute a placeholder node. Adds mapping between target of the node and the node itself to self.nodeDict.

Parameters:

target (Target) – The call target for this node. See Node for details on semantics
args (Tuple) – Tuple of positional args for this invocation
kwargs (Dict) – Dict of keyword arguments for this invocation

Returns:

A Proxy containing the node and the tracer of the new graph

transform() → dict

Run module via interpretation and return the the nodeDict.

Returns:: A dictionary containing the mapping between the target of the node and the node itself.