<filename>firewood/layers/conv_blocks.py
from typing import Any, Dict, List, Optional, Tuple, Union

from torch.nn.modules.utils import _pair, _single, _triple

from firewood.common.types import INT, SAME_PADDING
from firewood.layers.block import Block
from firewood.layers.conv_gradfix import (
    GFixConv1d,
    GFixConv2d,
    GFixConv3d,
    GFixConvTranspose1d,
    GFixConvTranspose2d,
    GFixConvTranspose3d,
)
from firewood.layers.conv_separable import (
    DepthSepConv1d,
    DepthSepConv2d,
    DepthSepConv3d,
    DepthSepConvTranspose1d,
    DepthSepConvTranspose2d,
    DepthSepConvTranspose3d,
    SpatialSepConv2d,
    SpatialSepConv3d,
    SpatialSepConvTranspose2d,
    SpatialSepConvTranspose3d,
)


class _ConvNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 1:
            convolution = GFixConv1d
        elif len(kernel_size) == 2:
            convolution = GFixConv2d  # type: ignore
        elif len(kernel_size) == 3:
            convolution = GFixConv3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class Conv1dBlock(_ConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _single(1),
        padding: Union[SAME_PADDING, INT] = _single(0),
        dilation: INT = _single(1),
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_single(kernel_size),
            stride=_single(stride),
            padding=padding,
            dilation=_single(dilation),
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class Conv2dBlock(_ConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class Conv3dBlock(_ConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class _ConvTransposeNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        output_padding: Tuple[int, ...],
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 1:
            convolution = GFixConvTranspose1d
        elif len(kernel_size) == 2:
            convolution = GFixConvTranspose2d  # type: ignore
        elif len(kernel_size) == 3:
            convolution = GFixConvTranspose3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            output_padding=output_padding,
            groups=groups,
            bias=bias,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class ConvTranspose1dBlock(_ConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _single(1),
        padding: Union[SAME_PADDING, INT] = _single(0),
        dilation: INT = _single(1),
        output_padding: INT = _single(0),
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_single(kernel_size),
            stride=_single(stride),
            padding=padding,
            dilation=_single(dilation),
            output_padding=_single(output_padding),
            groups=groups,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class ConvTranspose2dBlock(_ConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        output_padding: INT = _pair(0),
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            output_padding=_pair(output_padding),
            groups=groups,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class ConvTranspose3dBlock(_ConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        output_padding: INT = _triple(0),
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            output_padding=_triple(output_padding),
            groups=groups,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class _DepthSepConvNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        hidden_channels: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 1:
            convolution = DepthSepConv1d
        elif len(kernel_size) == 2:
            convolution = DepthSepConv2d  # type: ignore
        elif len(kernel_size) == 3:
            convolution = DepthSepConv3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            hidden_channels=hidden_channels,
            bias=bias,
            padding_mode=padding_mode,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class DepthSepConv1dBlock(_DepthSepConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _single(1),
        padding: Union[SAME_PADDING, INT] = _single(0),
        dilation: INT = _single(1),
        hidden_channels: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_single(kernel_size),
            stride=_single(stride),
            padding=padding,
            dilation=_single(dilation),
            hidden_channels=hidden_channels,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class DepthSepConv2dBlock(_DepthSepConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        hidden_channels: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            hidden_channels=hidden_channels,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class DepthSepConv3dBlock(_DepthSepConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        hidden_channels: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            hidden_channels=hidden_channels,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class _DepthSepConvTransposeNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        output_padding: Tuple[int, ...],
        hidden_channels: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 1:
            convolution = DepthSepConvTranspose1d
        elif len(kernel_size) == 2:
            convolution = DepthSepConvTranspose2d  # type: ignore
        elif len(kernel_size) == 3:
            convolution = DepthSepConvTranspose3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            output_padding=output_padding,
            hidden_channels=hidden_channels,
            bias=bias,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class DepthSepConvTranspose1dBlock(_DepthSepConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _single(1),
        padding: Union[SAME_PADDING, INT] = _single(0),
        dilation: INT = _single(1),
        output_padding: INT = _single(0),
        hidden_channels: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_single(kernel_size),
            stride=_single(stride),
            padding=padding,
            dilation=_single(dilation),
            output_padding=_single(output_padding),
            hidden_channels=hidden_channels,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class DepthSepConvTranspose2dBlock(_DepthSepConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        output_padding: INT = _pair(0),
        hidden_channels: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            output_padding=_pair(output_padding),
            hidden_channels=hidden_channels,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class DepthSepConvTranspose3dBlock(_DepthSepConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        output_padding: INT = _triple(0),
        hidden_channels: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            output_padding=_triple(output_padding),
            hidden_channels=hidden_channels,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class _SpatialSepConvNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 2:
            convolution = SpatialSepConv2d
        elif len(kernel_size) == 3:
            convolution = SpatialSepConv3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class SpatialSepConv2dBlock(_SpatialSepConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class SpatialSepConv3dBlock(_SpatialSepConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class _SpatialSepConvTransposeNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        output_padding: Tuple[int, ...],
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 2:
            convolution = SpatialSepConvTranspose2d
        elif len(kernel_size) == 3:
            convolution = SpatialSepConvTranspose3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            output_padding=output_padding,
            groups=groups,
            bias=bias,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class SpatialSepConvTranspose2dBlock(_SpatialSepConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        output_padding: INT = _pair(0),
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            output_padding=_pair(output_padding),
            groups=groups,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class SpatialSepConvTranspose3dBlock(_SpatialSepConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        output_padding: INT = _triple(0),
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            output_padding=_triple(output_padding),
            groups=groups,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )

<filename>firewood/layers/conv_blocks.py
from typing import Any, Dict, List, Optional, Tuple, Union

from torch.nn.modules.utils import _pair, _single, _triple

from firewood.common.types import INT, SAME_PADDING
from firewood.layers.block import Block
from firewood.layers.conv_gradfix import (
    GFixConv1d,
    GFixConv2d,
    GFixConv3d,
    GFixConvTranspose1d,
    GFixConvTranspose2d,
    GFixConvTranspose3d,
)
from firewood.layers.conv_separable import (
    DepthSepConv1d,
    DepthSepConv2d,
    DepthSepConv3d,
    DepthSepConvTranspose1d,
    DepthSepConvTranspose2d,
    DepthSepConvTranspose3d,
    SpatialSepConv2d,
    SpatialSepConv3d,
    SpatialSepConvTranspose2d,
    SpatialSepConvTranspose3d,
)


class _ConvNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 1:
            convolution = GFixConv1d
        elif len(kernel_size) == 2:
            convolution = GFixConv2d  # type: ignore
        elif len(kernel_size) == 3:
            convolution = GFixConv3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class Conv1dBlock(_ConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _single(1),
        padding: Union[SAME_PADDING, INT] = _single(0),
        dilation: INT = _single(1),
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_single(kernel_size),
            stride=_single(stride),
            padding=padding,
            dilation=_single(dilation),
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class Conv2dBlock(_ConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class Conv3dBlock(_ConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class _ConvTransposeNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        output_padding: Tuple[int, ...],
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 1:
            convolution = GFixConvTranspose1d
        elif len(kernel_size) == 2:
            convolution = GFixConvTranspose2d  # type: ignore
        elif len(kernel_size) == 3:
            convolution = GFixConvTranspose3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            output_padding=output_padding,
            groups=groups,
            bias=bias,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class ConvTranspose1dBlock(_ConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _single(1),
        padding: Union[SAME_PADDING, INT] = _single(0),
        dilation: INT = _single(1),
        output_padding: INT = _single(0),
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_single(kernel_size),
            stride=_single(stride),
            padding=padding,
            dilation=_single(dilation),
            output_padding=_single(output_padding),
            groups=groups,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class ConvTranspose2dBlock(_ConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        output_padding: INT = _pair(0),
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            output_padding=_pair(output_padding),
            groups=groups,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class ConvTranspose3dBlock(_ConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        output_padding: INT = _triple(0),
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            output_padding=_triple(output_padding),
            groups=groups,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class _DepthSepConvNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        hidden_channels: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 1:
            convolution = DepthSepConv1d
        elif len(kernel_size) == 2:
            convolution = DepthSepConv2d  # type: ignore
        elif len(kernel_size) == 3:
            convolution = DepthSepConv3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            hidden_channels=hidden_channels,
            bias=bias,
            padding_mode=padding_mode,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class DepthSepConv1dBlock(_DepthSepConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _single(1),
        padding: Union[SAME_PADDING, INT] = _single(0),
        dilation: INT = _single(1),
        hidden_channels: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_single(kernel_size),
            stride=_single(stride),
            padding=padding,
            dilation=_single(dilation),
            hidden_channels=hidden_channels,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class DepthSepConv2dBlock(_DepthSepConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        hidden_channels: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            hidden_channels=hidden_channels,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class DepthSepConv3dBlock(_DepthSepConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        hidden_channels: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            hidden_channels=hidden_channels,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class _DepthSepConvTransposeNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        output_padding: Tuple[int, ...],
        hidden_channels: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 1:
            convolution = DepthSepConvTranspose1d
        elif len(kernel_size) == 2:
            convolution = DepthSepConvTranspose2d  # type: ignore
        elif len(kernel_size) == 3:
            convolution = DepthSepConvTranspose3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            output_padding=output_padding,
            hidden_channels=hidden_channels,
            bias=bias,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class DepthSepConvTranspose1dBlock(_DepthSepConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _single(1),
        padding: Union[SAME_PADDING, INT] = _single(0),
        dilation: INT = _single(1),
        output_padding: INT = _single(0),
        hidden_channels: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_single(kernel_size),
            stride=_single(stride),
            padding=padding,
            dilation=_single(dilation),
            output_padding=_single(output_padding),
            hidden_channels=hidden_channels,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class DepthSepConvTranspose2dBlock(_DepthSepConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        output_padding: INT = _pair(0),
        hidden_channels: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            output_padding=_pair(output_padding),
            hidden_channels=hidden_channels,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class DepthSepConvTranspose3dBlock(_DepthSepConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        output_padding: INT = _triple(0),
        hidden_channels: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            output_padding=_triple(output_padding),
            hidden_channels=hidden_channels,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class _SpatialSepConvNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 2:
            convolution = SpatialSepConv2d
        elif len(kernel_size) == 3:
            convolution = SpatialSepConv3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class SpatialSepConv2dBlock(_SpatialSepConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class SpatialSepConv3dBlock(_SpatialSepConvNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            groups=groups,
            bias=bias,
            padding_mode=padding_mode,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class _SpatialSepConvTransposeNdBlock(Block):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Tuple[int, ...],
        stride: Tuple[int, ...],
        padding: Union[SAME_PADDING, INT],
        dilation: Tuple[int, ...],
        output_padding: Tuple[int, ...],
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        if len(kernel_size) == 2:
            convolution = SpatialSepConvTranspose2d
        elif len(kernel_size) == 3:
            convolution = SpatialSepConvTranspose3d  # type: ignore
        else:
            raise ValueError(f"Invalid kernel_size: {kernel_size}")
        weight_layer = convolution(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            output_padding=output_padding,
            groups=groups,
            bias=bias,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
        )
        super().__init__(
            weight_layer=weight_layer,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            dropout=dropout,
        )


class SpatialSepConvTranspose2dBlock(_SpatialSepConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _pair(1),
        padding: Union[SAME_PADDING, INT] = _pair(0),
        dilation: INT = _pair(1),
        output_padding: INT = _pair(0),
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        fir: Optional[List[float]] = None,
        fir_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_pair(kernel_size),
            stride=_pair(stride),
            padding=padding,
            dilation=_pair(dilation),
            output_padding=_pair(output_padding),
            groups=groups,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=fir,
            fir_args=fir_args,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )


class SpatialSepConvTranspose3dBlock(_SpatialSepConvTransposeNdBlock):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: INT,
        stride: INT = _triple(1),
        padding: Union[SAME_PADDING, INT] = _triple(0),
        dilation: INT = _triple(1),
        output_padding: INT = _triple(0),
        groups: int = 1,
        bias: bool = True,
        op_order: str = "WNA",
        normalization: Optional[str] = None,
        normalization_args: Optional[Dict[str, Any]] = None,
        activation: Optional[str] = None,
        activation_args: Optional[Dict[str, Any]] = None,
        noise: Optional[str] = None,
        noise_args: Optional[Dict[str, Any]] = None,
        weight_normalization: Optional[str] = None,
        weight_normalization_args: Optional[Dict[str, Any]] = None,
        lr_equalization: Optional[bool] = None,
        lr_equalization_args: Optional[Dict[str, Any]] = None,
        weight_initializer: str = "kaiming_uniform",
        bias_initializer: str = "zeros",
        dropout: Optional[float] = None,
    ) -> None:
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=_triple(kernel_size),
            stride=_triple(stride),
            padding=padding,
            dilation=_triple(dilation),
            output_padding=_triple(output_padding),
            groups=groups,
            bias=bias,
            op_order=op_order,
            normalization=normalization,
            normalization_args=normalization_args,
            activation=activation,
            activation_args=activation_args,
            fir=None,
            fir_args=None,
            noise=noise,
            noise_args=noise_args,
            weight_normalization=weight_normalization,
            weight_normalization_args=weight_normalization_args,
            lr_equalization=lr_equalization,
            lr_equalization_args=lr_equalization_args,
            weight_initializer=weight_initializer,
            bias_initializer=bias_initializer,
            dropout=dropout,
        )

# type: ignore # type: ignore # type: ignore # type: ignore # type: ignore # type: ignore # type: ignore # type: ignore # type: ignore # type: ignore