dmx.compressor.modeling.nn.torch_modules

Classes

`ActActMatMul`()
`AdaptiveAvgPool2d`(output_size)	An extension of PyTorch's AdaptiveAvgPool2d layer to support DmxModule configurations.
`AvgPool2d`(output_size)	An extension of PyTorch's AvgPool2d layer to support DmxModule configurations.
`BAddBMM`()
`BatchNorm2d`(num_features[, eps, momentum, ...])	An extension of PyTorch's BatchNorm2d layer to support DmxModule configurations.
`Conv1d`(in_channels, out_channels, kernel_size)	An extension of PyTorch's Conv1d layer to support DmxModule configurations.
`Conv2d`(in_channels, out_channels, kernel_size)	An extension of PyTorch's Conv2d layer to support DmxModule configurations.
`ConvTranspose2d`(in_channels, out_channels, ...)	An extension of PyTorch's ConvTranspose2d layer to support DmxModule configurations.
`Dropout`([p, inplace])	An extension of PyTorch's Dropout layer to support DmxModule configurations.
`Embedding`(num_embeddings, embedding_dim, ...)	An extension of PyTorch's Embedding layer to support DmxModule configurations.
`Exp`()
`GELU`(args, *kwargs)
`GELUBase`(activation_cls, args, *kwargs)	A generalized base class to support various GELUActivation configurations.
`GroupNorm`(num_groups, num_channels[, eps, ...])	An extension of PyTorch's GroupNorm layer to support DmxModule configurations.
`LayerNorm`(normalized_shape[, eps, ...])	An extension of PyTorch's LayerNorm layer to support DmxModule configurations.
`Linear`(in_features, out_features[, bias])	An extension of PyTorch's Linear layer to support DmxModule configurations.
`MaxPool2d`(kernel_size[, stride, padding, ...])	An extension of PyTorch's MaxPool2d layer to support DmxModule configurations.
`Mul`()
`RMSNorm`(normalized_shape[, eps])	An extension of RMSNorm layer to support DmxModule configurations.
`ReLU`([inplace])	An extension of PyTorch's ReLU layer to support DmxModule configurations.
`ReLU6`([inplace])	An extension of PyTorch's ReLU6 layer to support DmxModule configurations.
`ResAdd`()	A module for handling residual connections.
`ScaledDotProductAttention`([dropout_p])
`SiLU`([inplace])	An extension of PyTorch's SiLU (Sigmoid Linear Unit) layer to support DmxModule configurations.
`Softmax`([dim])	An extension of PyTorch's Softmax layer to support DmxModule configurations.
`Tanh`()	An extension of PyTorch's Tanh (Hyperbolic Tangent) layer to support DmxModule configurations.

class dmx.compressor.modeling.nn.torch_modules.ActActMatMul

Bases: DmxModule, Module

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.AdaptiveAvgPool2d(output_size)

Bases: DmxModule, AdaptiveAvgPool2d

An extension of PyTorch’s AdaptiveAvgPool2d layer to support DmxModule configurations. This module applies a 2D adaptive average pooling over an input signal composed of several input planes.

Parameters:: output_size (int or tuple) – The size of the output tensor after pooling.

None specific to this subclass. Inherits attributes from parent classes.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the 2D adaptive average pooling.

classmethod from_raw(raw: Module) → DmxModule

Creates a new AdaptiveAvgPool2d object (DmxModule) from a given PyTorch AdaptiveAvgPool2d layer.

Parameters:: raw (torch.nn.Module) – A PyTorch AdaptiveAvgPool2d layer to be converted.
Returns:: An AdaptiveAvgPool2d object that has the same configuration as the input PyTorch AdaptiveAvgPool2d layer.
Return type:: DmxModule

training: bool

class dmx.compressor.modeling.nn.torch_modules.AvgPool2d(output_size)

Bases: DmxModule, AvgPool2d

An extension of PyTorch’s AvgPool2d layer to support DmxModule configurations. This module applies a 2D average pooling over an input signal composed of several input planes.

Parameters:: output_size (int or tuple) – The size of the output tensor after pooling.

None specific to this subclass. Inherits attributes from parent classes.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the 2D average pooling.

classmethod from_raw(raw: Module) → DmxModule

Creates a new AvgPool2d object (DmxModule) from a given PyTorch AvgPool2d layer.

Parameters:: raw (torch.nn.Module) – A PyTorch AvgPool2d layer to be converted.
Returns:: An AvgPool2d object that has the same configuration as the input PyTorch AvgPool2d layer.
Return type:: DmxModule

training: bool

class dmx.compressor.modeling.nn.torch_modules.BAddBMM

Bases: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.BatchNorm2d(num_features: int, eps: float = 1e-05, momentum: float = 0.1, affine: bool = True, track_running_stats: bool = True)

Bases: DmxModule, BatchNorm2d

An extension of PyTorch’s BatchNorm2d layer to support DmxModule configurations. This module applies batch normalization over a 4D input tensor, suitable for use with 2D convolutional layers. The module is parameterized by the number of features, epsilon value for numerical stability, momentum for the running mean and variance, and options to use affine transformation and track running statistics.

Parameters:

num_features (int) – Number of channels in the input tensor.
eps (float, optional) – A small constant added to the denominator for numerical stability. Defaults to 1e-05.
momentum (float, optional) – The momentum value for the running mean and running variance computation. Defaults to 0.1.
affine (bool, optional) – Whether to include learnable affine parameters for this layer. Defaults to True.
track_running_stats (bool, optional) – Whether to track the running mean and variance during training. Defaults to True.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the batch normalization.

classmethod from_raw(raw: Module) → DmxModule

Creates a new BatchNorm2d object (DmxModule) from a given PyTorch BatchNorm2d layer.

Parameters:: raw (torch.nn.Module) – A PyTorch BatchNorm2d layer to be converted.
Returns:: A BatchNorm2d object that has the same configuration as the input PyTorch BatchNorm2d layer.
Return type:: DmxModule

training: bool

class dmx.compressor.modeling.nn.torch_modules.Conv1d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', **kwargs)

Bases: DmxModule, Conv1d

An extension of PyTorch’s Conv1d layer to support DmxModule configurations. This module performs a 1D convolution over an input signal composed of several input planes.

Parameters:

in_channels (int) – Number of channels in the input image.
out_channels (int) – Number of channels produced by the convolution.
kernel_size (int or tuple) – Size of the convolving kernel.
stride (int or tuple, optional) – Stride of the convolution. Defaults to 1.
padding (int or tuple, optional) – Zero-padding added to both sides of the input. Defaults to 0.
dilation (int or tuple, optional) – Spacing between kernel elements. Defaults to 1.
groups (int, optional) – Number of blocked connections from input channels to output channels. Defaults to 1.
bias (bool, optional) – If set to False, the layer will not learn an additive bias. Defaults to True.
padding_mode (str, optional) – Accepted values ‘zeros’ and ‘circular’ etc. Defaults to ‘zeros’.
**kwargs – Additional keyword arguments inherited from DmxModule.

_weight

The learnable weights of the module of shape (out_channels, in_channels, kernel_size).

Type:: Tensor

_bias

The learnable bias of the module of shape (out_channels).

Type:: Tensor, optional

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the 1D convolution.

classmethod from_raw(raw: Module) → DmxModule

Creates a new Conv1d object (DmxModule) from a given PyTorch Conv1d layer.

Parameters:: raw (torch.nn.Module) – A PyTorch Conv1d layer to be converted.
Returns:: A Conv1d object that has the same configuration as the input PyTorch Conv1d layer.
Return type:: DmxModule

training: bool

class dmx.compressor.modeling.nn.torch_modules.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', **kwargs)

Bases: DmxModule, Conv2d

An extension of PyTorch’s Conv2d layer to support DmxModule configurations. This module performs a 2D convolution over an input signal composed of several input planes.

Parameters:

in_channels (int) – Number of channels in the input image.
out_channels (int) – Number of channels produced by the convolution.
kernel_size (int or tuple) – Size of the convolving kernel.
stride (int or tuple, optional) – Stride of the convolution. Defaults to 1.
padding (int or tuple, optional) – Zero-padding added to both sides of the input. Defaults to 0.
dilation (int or tuple, optional) – Spacing between kernel elements. Defaults to 1.
groups (int, optional) – Number of blocked connections from input channels to output channels. Defaults to 1.
bias (bool, optional) – If set to False, the layer will not learn an additive bias. Defaults to True.
padding_mode (str, optional) – Accepted values ‘zeros’ and ‘circular’ etc. Defaults to ‘zeros’.
**kwargs – Additional keyword arguments inherited from DmxModule.

_weight

The learnable weights of the module of shape (out_channels, in_channels, kernel_height, kernel_width).

Type:: Tensor

_bias

The learnable bias of the module of shape (out_channels).

Type:: Tensor, optional

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the 2D convolution.

classmethod from_raw(raw: Module) → DmxModule

Creates a new Conv2d object (DmxModule) from a given PyTorch Conv2d layer.

Parameters:: raw (torch.nn.Module) – A PyTorch Conv2d layer to be converted.
Returns:: A Conv2d object that has the same configuration as the input PyTorch Conv2d layer.
Return type:: DmxModule

training: bool

class dmx.compressor.modeling.nn.torch_modules.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, padding_mode='zeros', **kwargs)

Bases: DmxModule, ConvTranspose2d

An extension of PyTorch’s ConvTranspose2d layer to support DmxModule configurations. This module performs a 2D transposed convolution over an input signal composed of several input planes.

Parameters:

in_channels (int) – Number of channels in the input image.
out_channels (int) – Number of channels produced by the transposed convolution.
kernel_size (int or tuple) – Size of the convolving kernel.
stride (int or tuple, optional) – Stride of the transposed convolution. Defaults to 1.
padding (int or tuple, optional) – Zero-padding added to both sides of the input. Defaults to 0.
output_padding (int or tuple, optional) – Additional size added to one side of the output shape. Defaults to 0.
groups (int, optional) – Number of blocked connections from input channels to output channels. Defaults to 1.
bias (bool, optional) – If set to False, the layer will not learn an additive bias. Defaults to True.
padding_mode (str, optional) – Accepted values ‘zeros’. Defaults to ‘zeros’.
**kwargs – Additional keyword arguments inherited from DmxModule.

_weight

The learnable weights of the module of shape (in_channels, out_channels, kernel_height, kernel_width).

Type:: Tensor

_bias

The learnable bias of the module of shape (out_channels).

Type:: Tensor, optional

_forward (_input: Tensor, output_size: Optional[List[int]] = None) -> Tensor: Computes the forward pass of the 2D transposed convolution.

classmethod from_raw(raw: Module) → DmxModule

training: bool

class dmx.compressor.modeling.nn.torch_modules.Dropout(p: float = 0.5, inplace: bool = False)

Bases: DmxModule, Dropout

An extension of PyTorch’s Dropout layer to support DmxModule configurations. This module applies the dropout operation over the input tensor.

Parameters:

p (float, optional) – The probability of an element to be zeroed. Defaults to 0.5.
inplace (bool, optional) – If set to True, will do this operation in-place. Defaults to False.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the dropout layer.

classmethod from_raw(raw: Module) → DmxModule

Creates a new Dropout object (DmxModule) from a given PyTorch Dropout layer.

Parameters:: raw (torch.nn.Module) – A PyTorch Dropout layer to be converted.
Returns:: A Dropout object that has the same configuration as the input PyTorch Dropout layer.
Return type:: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.Embedding(num_embeddings: int, embedding_dim: int, **kwargs)

Bases: DmxModule, Embedding

An extension of PyTorch’s Embedding layer to support DmxModule configurations.

Parameters:

num_embeddings (int) – size of the dictionary of embeddings
embedding_dim (int) – the size of each embedding vector
**kwargs – Additional keyword arguments inherited from torch.nn.Embedding.

_weight

the learnable weights of the module of shape (num_embeddings, embedding_dim).

Type:: Tensor

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the embedding layer.

classmethod from_raw(raw: Module) → DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.Exp

Bases: DmxModule, Module

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.GELU(*args, **kwargs)

Bases: GELUBase

functional_forward(x)

training: bool

class dmx.compressor.modeling.nn.torch_modules.GELUBase(activation_cls, *args, **kwargs)

Bases: DmxModule

A generalized base class to support various GELUActivation configurations. This module applies the specified GELUActivation function element-wise on the input tensor.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the GELU layer.

classmethod from_raw(raw: Module) → DmxModule

Creates a new GELU object (DmxModule) from a given Transformers layer.

Parameters:: raw (torch.nn.Module) – A Transformers GELUActivation layer to be converted.
Returns:: A GELU object that has the same configuration as the input Transformers GELUActivation layer.
Return type:: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.GroupNorm(num_groups: int, num_channels: int, eps: float = 1e-05, affine: bool = True)

Bases: DmxModule, GroupNorm

An extension of PyTorch’s GroupNorm layer to support DmxModule configurations. This module applies group normalization over an input tensor, suitable for use with various types of layers. The module is parameterized by the number of groups, number of channels, epsilon value for numerical stability, and an option to use affine transformation.

Parameters:

num_groups (int) – Number of groups to separate the channels into.
channels (int) – Number of channels in the input tensor.
eps (float, optional) – A small constant added to the denominator for numerical stability. Defaults to 1e-5.
affine (bool, optional) – Whether to include learnable affine parameters for this layer. Defaults to True.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the group normalization.

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.LayerNorm(normalized_shape: int | List[int] | Size, eps: float = 1e-05, elementwise_affine: bool = True)

Bases: DmxModule, LayerNorm

An extension of PyTorch’s LayerNorm layer to support DmxModule configurations. This module applies layer normalization over dimensions specified by the normalized_shape attribute. The mean and standard deviation are computed over the last D dimensions, where D`is the dimensionality indicated by `normalized_shape. Gamma and Beta are learnable parameters if elementwise_affine is set to True.

Parameters:

normalized_shape (Union[int, List[int], Size]) – Specifies dimensions for layer normalization.
eps (float, optional) – A value added for numerical stability. Defaults to 1e-5.
elementwise_affine (bool, optional) – Indicates if learnable affine parameters Gamma and Beta should be used. Defaults to True.

None specific to this subclass. Inherits attributes from parent classes.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the layer normalization.

approximator_wrapper(inputs, approx_args, approx_kwargs, **wrapper_kwargs): Override this in the DMX modules to enable pre-processing of the inputs, and the SIMD approximator arguments before calling the SIMD reference kernels. wrapper_kwargs are arguments specific to the wrapper and should not be passed to the SIMD reference kernels.

classmethod from_raw(raw: Module) → DmxModule

Creates a new LayerNorm object (DmxModule) from a given PyTorch LayerNorm layer.

Parameters:: raw (torch.nn.Module) – A PyTorch LayerNorm layer to be converted.
Returns:: A LayerNorm object that has the same configuration as the input PyTorch LayerNorm layer.
Return type:: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.Linear(in_features: int, out_features: int, bias: bool = True, **kwargs)

Bases: DmxModule, Linear

An extension of PyTorch’s Linear layer to support DmxModule configurations. This module performs a linear transformation on the input data.

Parameters:

in_features (int) – Size of each input sample.
out_features (int) – Size of each output sample.
bias (bool, optional) – If set to False, the layer will not learn an additive bias. Defaults to True.
**kwargs – Additional keyword arguments inherited from DmxModule.

_weight

The learnable weights of the module of shape (out_features, in_features).

Type:: Tensor

_bias

The learnable bias of the module of shape (out_features).

Type:: Tensor

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the linear layer.

classmethod from_raw(raw: Module) → DmxModule

to_compiler_graph() → DmxGraph

Returns a compiler friendly graph

>>> Reference:
    opcode         name                               target                             args                                                                               kwargs
    -------------  ---------------------------------  ---------------------------------  ---------------------------------------------------------------------------------  --------
    placeholder    _input                             _input                             ()                                                                                 {}
    get_attr       input_casts_input_cast_scale       input_casts.input_cast.scale       ()                                                                                 {}
    get_attr       input_casts_input_cast_zero_point  input_casts.input_cast.zero_point  ()                                                                                 {}
    call_function  quantize                           dmx.quantize                       (_input, input_casts_input_cast_scale, input_casts_input_cast_zero_point, 'SAME')  {}
    call_function  dequantize                         dmx.dequantize                     (quantize, input_casts_input_cast_scale, input_casts_input_cast_zero_point)        {}
    get_attr       _weight                            _weight                            ()                                                                                 {}
    get_attr       weight_scale                       weight_scale                       ()                                                                                 {}
    get_attr       weight_zero_point                  weight_zero_point                  ()                                                                                 {}
    call_function  quantize_1                         dmx.quantize                       (_weight, weight_scale, weight_zero_point, 'SAME')                                 {}
    call_function  dequantize_1                       dmx.dequantize                     (quantize_1, weight_scale, weight_zero_point)                                      {}
    get_attr       _bias                              _bias                              ()                                                                                 {}
    get_attr       bias_cast_scale                    bias_cast.scale                    ()                                                                                 {}
    get_attr       bias_cast_zero_point               bias_cast.zero_point               ()                                                                                 {}
    call_function  quantize_2                         dmx.quantize                       (_bias, bias_cast_scale, bias_cast_zero_point, 'SAME')                             {}
    call_function  dequantize_2                       dmx.dequantize                     (quantize_2, bias_cast_scale, bias_cast_zero_point)                                {}
    call_function  _output                            <built-in function linear>         (dequantize, dequantize_1, dequantize_2)                                           {}
    get_attr       output_casts_output_cast_scale     output_casts.output_cast.scale     ()                                                                                 {}
    get_attr       output_casts_output_cast_zero_pointoutput_casts.utput_cast.zero_point ()                                                                                 {}
    call_function  quantize_3                         dmx.quantize                       (_output, output_casts_output_cast_scale, output_casts_output_cast_zero_pointoutput_casts, 'SAME')  {}
    call_function  dequantize_3                       dmx.dequantize                     (quantize_3, output_casts_output_cast_scale, output_casts_output_cast_zero_pointoutput_casts)       {}
    output         output                             output                             (dequantize_3,)                                                                    {}

training: bool

class dmx.compressor.modeling.nn.torch_modules.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)

Bases: DmxModule, MaxPool2d

An extension of PyTorch’s MaxPool2d layer to support DmxModule configurations. This module applies a 2D max pooling over an input signal composed of several input planes.

Parameters:

kernel_size (int or tuple) – Size of the window to take a max over.
stride (int or tuple, optional) – Stride of the window. Defaults to None.
padding (int or tuple, optional) – Zero-padding added to both sides of the input. Defaults to 0.
dilation (int or tuple, optional) – Spacing between kernel elements. Defaults to 1.
return_indices (bool, optional) – If True, will return the max indices in a second tensor. Defaults to False.
ceil_mode (bool, optional) – If True, will use ceil instead of floor to compute the output shape. Defaults to False.

None specific to this subclass. Inherits attributes from parent classes.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the 2D max pooling.

classmethod from_raw(raw: Module) → DmxModule

Creates a new MaxPool2d object (DmxModule) from a given PyTorch MaxPool2d layer.

Parameters:: raw (torch.nn.Module) – A PyTorch MaxPool2d layer to be converted.
Returns:: A MaxPool2d object that has the same configuration as the input PyTorch MaxPool2d layer.
Return type:: DmxModule

training: bool

class dmx.compressor.modeling.nn.torch_modules.Mul

Bases: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.RMSNorm(normalized_shape: int, eps: float = 1e-06)

Bases: DmxModule, RMSNorm

An extension of RMSNorm layer to support DmxModule configurations. This module performs RMS-based layer normalization on the input tensor. The layer normalization is characterized by the hidden_size and an optional eps value for numerical stability.

Parameters:

dim (int) – The size of the hidden layer (number of hidden units).
eps (float, optional) – A small constant added to the denominator for numerical stability. Defaults to 1e-6.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the RMS layer normalization.

classmethod from_raw(raw: Module) → DmxModule

Creates a new RMSNorm object (DmxModule) from a given PyTorch RMSNorm layer.

Parameters:: raw (torch.nn.Module) – A PyTorch RMSNorm layer to be converted.
Returns:: A RMSNorm object that has the same configuration as the input PyTorch RMSNorm layer.
Return type:: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.ReLU(inplace: bool = False)

Bases: DmxModule, ReLU

An extension of PyTorch’s ReLU layer to support DmxModule configurations. This module applies the Rectified Linear Unit (ReLU) function element-wise on the input tensor.

Parameters:: inplace (bool, optional) – If set to True, will do this operation in-place. Defaults to False.

_forward (_input: Tensor, inplace: bool = False) -> Tensor: Computes the forward pass of the ReLU layer.

classmethod from_raw(raw: Module) → DmxModule

Creates a new ReLU object (DmxModule) from a given PyTorch ReLU layer.

Parameters:: raw (torch.nn.Module) – A PyTorch ReLU layer to be converted.
Returns:: A ReLU object that has the same configuration as the input PyTorch ReLU layer.
Return type:: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.ReLU6(inplace: bool = False)

Bases: DmxModule, ReLU6

An extension of PyTorch’s ReLU6 layer to support DmxModule configurations. This module applies the Rectified Linear Unit 6 (ReLU6) function element-wise on the input tensor.

Parameters:: inplace (bool, optional) – If set to True, will do this operation in-place. Defaults to False.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the ReLU6 layer.

classmethod from_raw(raw: Module) → DmxModule

Creates a new ReLU6 object (DmxModule) from a given PyTorch ReLU6 layer.

Parameters:: raw (torch.nn.Module) – A PyTorch ReLU6 layer to be converted.
Returns:: A ReLU6 object that has the same configuration as the input PyTorch ReLU6 layer.
Return type:: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.ResAdd

Bases: DmxModule, Module

A module for handling residual connections.

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.ScaledDotProductAttention(dropout_p=0.0)

Bases: DmxModule

forward(query, key, value, attn_mask=None, is_causal=False, scale=None, enable_gqa=False)

Forward pass of the module with quantization ops applied.

Parameters:

input (Tensor) – input tensor to be passed through the module
*args (Optional[Tuple]) – variable length of args
**kwargs (Optional[Dict]) – variable length of kwargs

is_compound = True

module_graph(*args, **kwargs) → Graph

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.SiLU(inplace: bool = False)

Bases: DmxModule, SiLU

An extension of PyTorch’s SiLU (Sigmoid Linear Unit) layer to support DmxModule configurations. This module applies the SiLU function element-wise on the input tensor.

Parameters:: inplace (bool, optional) – If set to True, will do this operation in-place. Defaults to False.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the SiLU layer.

classmethod from_raw(raw: Module) → DmxModule

Creates a new SiLU object (DmxModule) from a given PyTorch SiLU layer.

Parameters:: raw (torch.nn.Module) – A PyTorch SiLU layer to be converted.
Returns:: A SiLU object that has the same configuration as the input PyTorch SiLU layer.
Return type:: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.Softmax(dim: int = -1)

Bases: DmxModule, Softmax

An extension of PyTorch’s Softmax layer to support DmxModule configurations. This module applies the Softmax function to an n-dimensional input tensor, normalizing the elements along a specified dimension such that they sum up to 1.

Parameters:: dim (int, optional) – Dimension along which Softmax will be computed. Defaults to -1.

None specific to this subclass. Inherits attributes from parent classes.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the Softmax function.

approximator_wrapper(inputs, approx_args, approx_kwargs, **wrapper_kwargs): Override this in the DMX modules to enable pre-processing of the inputs, and the SIMD approximator arguments before calling the SIMD reference kernels. wrapper_kwargs are arguments specific to the wrapper and should not be passed to the SIMD reference kernels.

classmethod from_raw(raw: Module) → DmxModule

Creates a Softmax DmxModule instance from a raw torch.nn.Module instance.

Parameters:: raw (torch.nn.Module) – The raw torch.nn.Module instance.
Returns:: An initialized Softmax DmxModule instance with parameters copied from the raw instance.
Return type:: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool

class dmx.compressor.modeling.nn.torch_modules.Tanh

Bases: DmxModule, Tanh

An extension of PyTorch’s Tanh (Hyperbolic Tangent) layer to support DmxModule configurations. This module applies the tanh function element-wise on the input tensor.

_forward (_input: Tensor) -> Tensor: Computes the forward pass of the Tanh layer.

classmethod from_raw(raw: Module) → DmxModule

Creates a new Tanh object (DmxModule) from a given PyTorch Tanh layer.

Parameters:: raw (torch.nn.Module) – A PyTorch Tanh layer to be converted.
Returns:: A Tanh object that has the same configuration as the input PyTorch Tanh layer.
Return type:: DmxModule

to_compiler_graph() → Graph: Returns a compiler friendly graph

training: bool