custom_layers

def __init__(
    self,
    in_features: int,
    nonlin: Union[ActivationFunction, Sequence[ActivationFunction]],
    bias: bool,
    weight: bool = True,
):
    """
    Args:
        in_features: Number of dimensions of each input sample.
        nonlin: The nonlinear function to apply.
        bias: If `True`, a learnable bias is subtracted, else no bias is used.
        weight: If `True`, the input is multiplied with a learnable scaling factor, else no weighting is used.
    """
    if not callable(nonlin):
        if len(nonlin) != in_features:
            raise RuntimeError(
                f"Either one, or {in_features} nonlinear functions have been expected, but "
                f"{len(nonlin)} have been given!"
            )

    super().__init__()

    # Create and initialize the parameters, and the activation function.
    self.nonlin = copy.deepcopy(nonlin) if _is_iterable(nonlin) else nonlin
    if weight:
        self.weight = nn.Parameter(torch.empty(in_features, dtype=torch.get_default_dtype()))
    else:
        self.register_buffer("weight", torch.ones(in_features, dtype=torch.get_default_dtype()))
    if bias:
        self.bias = nn.Parameter(torch.empty(in_features, dtype=torch.get_default_dtype()))
    else:
        self.register_buffer("bias", torch.zeros(in_features, dtype=torch.get_default_dtype()))

    init_param_(self)

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    stride: int = 1,
    padding: Union[int, str] = "same",  # kernel_size // 2 if padding_mode != "circular" else kernel_size - 1
    dilation: int = 1,
    groups: int = 1,
    bias: bool = False,
    padding_mode: str = "zeros",
    device: Optional[Union[str, torch.device]] = None,
    dtype=None,
):
    # Same as in PyTorch 1.12.
    super().__init__(
        in_channels=in_channels,
        out_channels=out_channels,
        kernel_size=_single(kernel_size),  # type: ignore[arg-type]
        stride=_single(stride),  # type: ignore[arg-type]
        padding=_single(padding) if not isinstance(padding, str) else padding,  # type: ignore[arg-type]
        dilation=_single(dilation),  # type: ignore[arg-type]
        transposed=False,
        output_padding=_single(0),
        groups=groups,
        bias=bias,
        padding_mode=padding_mode,
        device=device,
        dtype=dtype,
    )

    # Memorize PyTorch's weight shape (out_channels x in_channels x kernel_size) for later reconstruction.
    self.orig_weight_shape = self.weight.shape

    # Get number of kernel elements we later want to use for mirroring.
    self.half_kernel_size = math.ceil(self.weight.size(2) / 2)  # kernel_size = 4 --> 2, kernel_size = 5 --> 3

    # Initialize the weights values the same way PyTorch does.
    new_weight_init = torch.zeros(
        self.orig_weight_shape[0], self.orig_weight_shape[1], self.half_kernel_size, device=device
    )
    nn.init.kaiming_uniform_(new_weight_init, a=math.sqrt(5))

    # Overwrite the weight attribute (transposed is False by default for the Conv1d module, we don't use it here).
    self.weight = nn.Parameter(new_weight_init)