dmx.compressor.fx.transformer.dmx_aware_transformer

Classes

DMXAwareTransformer(module, node_name_to_scope)

Substitute as in dmx.model.aware(), replace torch.nn.modules and activations with dmx counterpart
class dmx.compressor.fx.transformer.dmx_aware_transformer.DMXAwareTransformer(module: GraphModule, node_name_to_scope: dict, old_gms: List[GraphModule] = None, nodeInputs: dict = None)

Bases: Transformer

Substitute as in dmx.model.aware(), replace torch.nn.modules and activations with dmx counterpart

Parameters:

  • module (fx.GraphModule) – the module to transform

  • node_name_to_scope (dict) – A dictionary storing the mapping between node names and scopes

module

the module to transform

Type:

fx.GraphModule

node_name_to_scope

A dictionary storing the mapping between node names and scopes

Type:

dict

old_gms List[

dmxmodules to reuse if module was already transformed

Type:

fx.GraphModule

add_dmx_aware_functional_mapping(target: str, dmx_module_cls)
add_submod(name, mod)

this function will try to reuse modules in old_gm if possible

call_function(target: Callable[[...], Any] | str, args: Tuple[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, ...] | Sequence[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] | Mapping[str, 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] | 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, replaces the function with its dmx counterpart and returns the result.

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 new node and the tracer of the new graph

call_method(target: Callable[[...], Any] | str, args: Tuple[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, ...] | Sequence[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] | Mapping[str, 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] | 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 and return the result.

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

Return

Any: The value returned by the method invocation

Note

Backwards-compatibility for this API is guaranteed.

call_module(target: Callable[[...], Any] | str, args: Tuple[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, ...] | Sequence[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] | Mapping[str, 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] | 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, replaces the module with its dmx counterpart and returns the result.

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 new node and the tracer of the new graph

create_unique_name_in_scope(cand_name)
get_attr(target, args, kwargs)

Execute a get_attr node. In Transformer, this is overridden to insert a new get_attr node into the output graph.

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

Note

Backwards-compatibility for this API is guaranteed.

merge_graph(subgraph, scope, args, kwargs)

This function is to merge subgraph to graph

placeholder(target, args, kwargs)

Execute a placeholder node. In Transformer, this is overridden to insert a new placeholder into the output graph.

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

Note

Backwards-compatibility for this API is guaranteed.