# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/blocks/04_cnn.ipynb.

# %% auto 0
__all__ = ['possible_sequences', 'CNNBlockConfig', 'MultiResCNNBlockConfig', 'CNNBlock3D', 'CNNBlock2D', 'MultiResCNNBlock3D',
           'MultiResCNNBlock2D', 'TensorSplittingConv', 'add_tsp_to_module', 'remove_tsp_from_module']

# %% ../../nbs/blocks/04_cnn.ipynb 2
from functools import cache, wraps
from itertools import chain, permutations
from typing import Any, Literal

import torch
from loguru import logger
from torch import nn
from torch.nn import functional as F

from ..docstrings import populate_docstring
from ..utils.activation_checkpointing import ActivationCheckpointing
from ..utils.activations import get_act_layer
from ..utils.custom_base_model import CustomBaseModel, Field, field_validator, model_validator
from ..utils.normalizations import get_norm_layer
from ..utils.rearrange import rearrange_channels
from ..utils.residuals import Residual
from ..utils.splitter_merger import Splitter

# %% ../../nbs/blocks/04_cnn.ipynb 4
possible_sequences = ["".join(p) for p in chain.from_iterable(permutations("ACDN", r) for r in range(5)) if "C" in p]


class CNNBlockConfig(CustomBaseModel):
    in_channels: int = Field(..., description="Number of input channels")
    out_channels: int = Field(..., description="Number of output channels")
    kernel_size: int | tuple[int, ...] = Field(..., description="Kernel size for the convolution")
    padding: int | tuple[int, ...] | str = Field(
        "same", description="Padding for the convolution. Can be 'same' or an integer/tuple of integers."
    )
    stride: int | tuple[int, ...] = Field(1, description="Stride for the convolution")
    conv_kwargs: dict[str, Any] = Field({}, description="Additional keyword arguments for the convolution layer")
    transposed: bool = Field(False, description="Whether to perform ConvTranspose instead of Conv")

    normalization: str | None = Field("batchnorm3d", description="Normalization layer type.")
    normalization_pre_args: list = Field(
        [],
        description=(
            "Arguments for the normalization layer before providing the dimension. Useful when using "
            "GroupNorm layers are being used to specify the number of groups."
        ),
    )
    normalization_post_args: list = Field(
        [], description="Arguments for the normalization layer after providing the dimension."
    )
    normalization_kwargs: dict = Field({}, description="Additional keyword arguments for the normalization layer")
    activation: str | None = Field("relu", description="Activation function type.")
    activation_kwargs: dict = Field({}, description="Additional keyword arguments for the activation function.")

    sequence: Literal[tuple(possible_sequences)] = Field("CNA", description="Sequence of operations in the block.")

    drop_prob: float = Field(0.0, description="Dropout probability.")

    @model_validator(mode="after")
    def validate(self):
        super().validate()
        if self.normalization is None and "N" in self.sequence:
            self.sequence = self.sequence.replace("N", "")
        if self.normalization is not None and "N" not in self.sequence:
            raise ValueError("Add N to the sequence or set normalization=None.")
        if self.activation is None and "A" in self.sequence:
            self.sequence = self.sequence.replace("A", "")
        if self.activation is not None and "A" not in self.sequence:
            raise ValueError("Add A to the sequence or set activation=None.")
        if self.drop_prob == 0.0 and "D" in self.sequence:
            self.sequence = self.sequence.replace("D", "")
        if self.drop_prob > 0.0 and "D" not in self.sequence:
            raise ValueError("Add D to the sequence or set drop_prob=0.")
        return self


class MultiResCNNBlockConfig(CNNBlockConfig):
    kernel_sizes: tuple[int | tuple[int, ...], ...] = Field((3, 5, 7), description="Kernel sizes for each conv layer.")
    filter_ratios: tuple[float, ...] = Field(
        (1, 2, 3), description="Ratio of filters to out_channels for each conv layer. Will be scaled to sum to 1."
    )
    padding: Literal["same"] = Field(
        "same", description="Padding for the convolution. Only 'same' is supported for MultiResCNNBlock."
    )

    kernel_size: int = Field(
        3, description="Kernel size for the convolution. Only kernel_size=3 is supported for MultiResCNNBlock."
    )

    @field_validator("filter_ratios", mode="after")
    @classmethod
    def scale_filter_ratios(cls, filter_ratios):
        filter_ratios = tuple(ratio / sum(filter_ratios) for ratio in filter_ratios)
        return filter_ratios

    @model_validator(mode="after")
    def validate(self):
        super().validate()
        assert self.kernel_sizes == (3, 5, 7), "Only kernel sizes of (3, 5, 7) are supported for MultiResCNNBlock"
        assert self.kernel_size == 3, "only kernel_size = 3 is supported for MultiResCNNBlock"
        assert len(self.kernel_sizes) == len(
            self.filter_ratios
        ), "kernel_sizes and filter_ratios must have the same length"
        return self

# %% ../../nbs/blocks/04_cnn.ipynb 7
@populate_docstring
class _CNNBlock(nn.Module):
    """A block to perform a sequence of convolution, activation, normalization, and dropout operations. Works for 2D as
    well as 3D data."""

    @populate_docstring
    def __init__(
        self, spatial_dims: Literal[2, 3], config: CNNBlockConfig = {}, checkpointing_level: int = 0, **kwargs
    ):
        """Initialize the CNNBlock block. Activation checkpointing level 1.

        Args:
            spatial_dims: Number of spatial dimensions (2 or 3). This is used to determine 2D vs 3D data
            config: {CONFIG_INSTANCE_DOC}
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__()

        self.config = CNNBlockConfig.model_validate(config | kwargs)

        normalization = self.config.normalization
        activation = self.config.activation
        drop_prob = self.config.drop_prob
        sequence = self.config.sequence

        bias = True
        if normalization is not None and normalization.startswith("batchnorm") and "CN" in sequence:
            bias = False

        match spatial_dims, self.config.transposed:
            case 2, False:
                conv_module = nn.Conv2d
            case 2, True:
                conv_module = nn.ConvTranspose2d
            case 3, False:
                conv_module = nn.Conv3d
            case 3, True:
                conv_module = nn.ConvTranspose3d
            case _:
                raise ValueError(f"Unsupported spatial dimensions: {spatial_dims}")

        self.conv = conv_module(
            in_channels=self.config.in_channels,
            out_channels=self.config.out_channels,
            kernel_size=self.config.kernel_size,
            padding=self.config.padding,
            stride=self.config.stride,
            bias=bias,
            **self.config.conv_kwargs,
        )

        self.norm = None
        self.act = None
        self.dropout = None

        norm_channels = self.config.out_channels
        if "N" in sequence.split("C")[0]:
            norm_channels = self.config.in_channels

        if "N" in sequence:
            self.norm = get_norm_layer(
                normalization,
                *self.config.normalization_pre_args,
                norm_channels,
                *self.config.normalization_post_args,
                **self.config.normalization_kwargs,
            )
        if "A" in sequence:
            self.act = get_act_layer(activation, **self.config.activation_kwargs)
        if "D" in sequence:
            self.dropout = nn.Dropout(drop_prob)

        self.checkpointing_level1 = ActivationCheckpointing(1, checkpointing_level)

    @populate_docstring
    def _forward(self, x: torch.Tensor, channels_first: bool = True) -> torch.Tensor:
        """Forward pass of the CNNBlock block.

        Args:
            x: {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            {OUTPUT_3D_DOC}
        """
        # x: (b, [in_channels], [z], y, x, [in_channels])

        x = rearrange_channels(x, channels_first, True)
        # Now x is (b, in_channels, [z], y, x)

        for layer in self.config.sequence:
            if layer == "C":
                x = self.conv(x)
            if layer == "A":
                x = self.act(x)
            elif layer == "D":
                x = self.dropout(x)
            elif layer == "N":
                x = self.norm(x)
        # (b, out_channels, [z], y, x)

        x = rearrange_channels(x, True, channels_first)
        # (b, [out_channels], [z], y, x, [out_channels])

        return x

    @wraps(_forward)
    def forward(self, *args, **kwargs):
        return self.checkpointing_level1(self._forward, *args, **kwargs)

# %% ../../nbs/blocks/04_cnn.ipynb 8
@populate_docstring
class CNNBlock3D(_CNNBlock):
    """A block to perform a sequence of convolution, activation, normalization, and dropout operations.
    {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: CNNBlockConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the CNNBlock3D block. Activation checkpointing level 1.

        Args:
            config: {CONFIG_INSTANCE_DOC}
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__(3, config, checkpointing_level, **kwargs)

# %% ../../nbs/blocks/04_cnn.ipynb 11
@populate_docstring
class CNNBlock2D(_CNNBlock):
    """A block to perform a sequence of convolution, activation, normalization, and dropout operations.
    {CLASS_DESCRIPTION_2D_DOC}"""

    @populate_docstring
    def __init__(self, config: CNNBlockConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the CNNBlock2D block. Activation checkpointing level 1.

        Args:
            config: {CONFIG_INSTANCE_DOC}
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__(2, config, checkpointing_level, **kwargs)

# %% ../../nbs/blocks/04_cnn.ipynb 15
class _MultiResCNNBlock(nn.Module):
    """A block to perform a multi-resolution convolution operation using the cascading convolutions trick. It uses three
    different kernel sizes (3, 5, and 7) to capture features at multiple resolutions. The output channels are split
    according to the specified filter ratios. The block also includes a residual connection."""

    @populate_docstring
    def __init__(
        self, spatial_dims: Literal[2, 3], config: MultiResCNNBlockConfig = {}, checkpointing_level: int = 0, **kwargs
    ):
        """Initialize the MultiResCNNBlock block. Activation checkpointing level 2.

        Args:
            spatial_dims: Number of spatial dimensions (2 or 3). This is used to determine 2D vs 3D data
            config: {CONFIG_INSTANCE_DOC}
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__()

        self.config = MultiResCNNBlockConfig.model_validate(config | kwargs)

        assert self.config.kernel_sizes == (3, 5, 7), "Only kernel sizes of (3, 5, 7) are supported for now"

        all_out_channels = [max(1, int(self.config.out_channels * ratio)) for ratio in self.config.filter_ratios[:-1]]
        last_out_channels = self.config.out_channels - sum(all_out_channels)
        all_out_channels.append(last_out_channels)
        if last_out_channels <= 0:
            raise ValueError(
                f"These filter values ({self.config.filter_ratios}) won't work with the given out_channels. Please "
                f"adjust them. The out_channels of each conv layer is coming out to be {all_out_channels}."
            )
        all_in_channels = [self.config.in_channels] + all_out_channels[:-1]

        self.convs = nn.ModuleList(
            [
                _CNNBlock(
                    spatial_dims,
                    self.config.model_dump(),
                    checkpointing_level,
                    in_channels=in_channels,
                    out_channels=out_channels,
                    kernel_size=3,
                )
                for in_channels, out_channels in zip(all_in_channels, all_out_channels)
            ]
        )

        self.residual_conv = _CNNBlock(
            spatial_dims,
            self.config.model_dump(),
            checkpointing_level,
            in_channels=self.config.in_channels,
            out_channels=self.config.out_channels,
            kernel_size=1,
        )

        self.residual = Residual()

        self.checkpointing_level2 = ActivationCheckpointing(2, checkpointing_level)

    @populate_docstring
    def _forward(self, x: torch.Tensor, channels_first: bool = True) -> torch.Tensor:
        """Forward pass of the MultiResCNNBlock block.

        Args:
            x: {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            {OUTPUT_3D_DOC}
        """
        # x: (b, [in_channels], [z], y, x, [in_channels])

        x = rearrange_channels(x, channels_first, True)
        # (b, in_channels, [z], y, x)

        residual = self.residual_conv(x)
        # (b, out_channels, [z], y, x)

        conv_outputs = []
        for conv in self.convs:
            conv_input = conv_outputs[-1] if conv_outputs else x
            conv_output = conv(conv_input)
            conv_outputs.append(conv_output)
            # (b, one_of_all_out_channels, [z], y, x)

        x = torch.cat(conv_outputs, dim=1)
        # (b, out_channels, [z], y, x)

        x = self.residual(x, residual)
        # (b, out_channels, [z], y, x)

        x = rearrange_channels(x, True, channels_first)
        # (b, [out_channels], [z], y, x, [out_channels])

        return x

    @wraps(_forward)
    def forward(self, *args, **kwargs):
        return self.checkpointing_level2(self._forward, *args, **kwargs)

# %% ../../nbs/blocks/04_cnn.ipynb 16
class MultiResCNNBlock3D(_MultiResCNNBlock):
    @populate_docstring
    def __init__(self, config: MultiResCNNBlockConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the MultiResCNNBlock3D block. Activation checkpointing level 2.

        Args:
            config: {CONFIG_INSTANCE_DOC}
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__(3, config, checkpointing_level, **kwargs)

# %% ../../nbs/blocks/04_cnn.ipynb 18
class MultiResCNNBlock2D(_MultiResCNNBlock):
    @populate_docstring
    def __init__(self, config: MultiResCNNBlockConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the MultiResCNNBlock2D block. Activation checkpointing level 2.

        Args:
            config: {CONFIG_INSTANCE_DOC}
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__(2, config, checkpointing_level, **kwargs)

# %% ../../nbs/blocks/04_cnn.ipynb 21
class TensorSplittingConv(nn.Module):
    """Convolution layer that operates on splits of a tensor on desired device and concatenates the results to give a
    lossless output. This is useful for large input tensors that cause intermediate buffers in the conv layer that
    don't fit in memory. Works for both 2D and 3D convolutions."""

    def __init__(self, conv: nn.Module, num_splits: int | tuple[int, ...], optimize_num_splits: bool = True):
        """Initialize the TensorSplittingConv layer.

        Args:
            conv: Convolution layer to be used for splitting. Must be either nn.Conv2d or nn.Conv3d.
            num_splits: Number of splits for each spatial dimension. If an int is provided, it will be used for all
                spatial dimensions. If a tuple is provided, it must have the same length as the number of spatial
                dimensions.
            optimize_num_splits: Whether to optimize the number of splits based on the input shape. An example of
                optimization is provided below. Defaults to True.
        """
        super().__init__()

        if isinstance(conv, nn.Conv2d):
            self.spatial_dims = 2
        elif isinstance(conv, nn.Conv3d):
            self.spatial_dims = 3
        else:
            raise ValueError("Unsupported convolution type. Only Conv2d and Conv3d are supported.")

        assert conv.stride == (1,) * self.spatial_dims, "Stride must be 1 for tensor splitting convolution."
        assert conv.padding == "same" or torch.allclose(
            torch.tensor(conv.padding), (torch.tensor(conv.kernel_size) - 1) // 2
        ), "Padding must be 'same' for tensor splitting convolution."

        if isinstance(num_splits, int):
            num_splits = (num_splits,) * self.spatial_dims
        assert len(num_splits) == self.spatial_dims, "num_splits must be a tuple of length equal to spatial_dims"

        self.conv = conv
        self.num_splits = num_splits
        self.optimize_num_splits = optimize_num_splits

    @cache
    def get_receptive_field(self) -> tuple[int, ...]:
        """Calculate the receptive field of the convolution layer."""
        kernel_size = torch.tensor(self.conv.kernel_size)
        dilation = torch.tensor(self.conv.dilation)
        receptive_field = dilation * (kernel_size - 1) + 1
        return tuple(receptive_field.tolist())

    @cache
    def get_edge_context(self):
        """Calculate the context size required to eliminate edge effects when merging the conv outputs into one."""
        receptive_field = self.get_receptive_field()
        context = torch.tensor(receptive_field) // 2
        return tuple(context.tolist())

    def get_input_shape(self, input_shape: tuple[int, ...] | torch.Size | torch.Tensor) -> tuple[int, ...]:
        """Get the input shape of the convolution layer. This function removes any unnecesary dimensions and ensures
        that the input shape is of length equal to the number of spatial dimensions.

        Args:
            input_shape: Shape of the input tensor. If a tensor is provided, its shape will be used.

        Returns:
            Tuple of the input shape for the convolution layer, with only the spatial dimensions.
        """
        if isinstance(input_shape, torch.Tensor):
            input_shape = input_shape.shape
        if isinstance(input_shape, torch.Size):
            input_shape = tuple(input_shape)
        input_shape = input_shape[-self.spatial_dims :]
        if len(input_shape) != self.spatial_dims:
            raise ValueError(f"Input shape must be of length {self.spatial_dims}. Got {len(input_shape)}.")
        return input_shape

    def get_optimized_num_splits(self, input_shape: tuple[int, ...] | torch.Size | torch.Tensor) -> tuple[int, ...]:
        """Optimize the number of splits for each dimension based on the input shape and number of splits

        Example:
            Let's say input shape is (110, 110) and num_splits is (12, 12). The input will first be padded to (120, 120)
            and then split into splits of size (10+overlap, 10+overlap) each. However, if you notice, the padding
            that was also equal to 10, and therefore was completely unnecessary as the same result can be achieved by
            using num_splits = (11, 11) and reducing 144-121=23 splits to be processed.

        Args:
            input_shape: Shape of the input tensor. If a tensor is provided, its shape will be used.

        Returns:
            Tuple of optimized number of splits for each dimension.
        """
        input_shape = self.get_input_shape(input_shape)

        num_splits = list(self.num_splits)
        for i in range(self.spatial_dims):
            while True:
                padding_required = (num_splits[i] - (input_shape[i] % num_splits[i])) % num_splits[i]
                split_size = (input_shape[i] + padding_required) // num_splits[i]
                if padding_required >= split_size:
                    num_splits[i] -= 1
                else:
                    break

        return tuple(num_splits)

    def pad_input_for_divisibility(self, x: torch.Tensor, num_splits: tuple[int, ...] = None) -> torch.Tensor:
        """Pad the input at the end of every spatial dimension such that it is perfectly divisible by the number of
        splits.

        Args:
            x: Input tensor of shape (batch_size, in_channels, [z], y, x).
            num_splits: Number of splits for each spatial dimension. If None, the default num_splits will be used.

        Returns:
            Padded input tensor of shape (batch_size, in_channels, [z], y, x).
        """
        if num_splits is None:
            num_splits = self.num_splits
        padding = [0, 0] * (x.ndim - self.spatial_dims)
        for i in range(self.spatial_dims):
            dim = i + 2
            padding_required = (num_splits[i] - (x.shape[dim] % num_splits[i])) % num_splits[i]
            padding.extend([padding_required, 0])
        x = F.pad(x, list(reversed(padding)))
        return x

    def get_split_size(
        self, input_shape: tuple[int, ...] | torch.Size | torch.Tensor, num_splits: tuple[int, ...] = None
    ) -> tuple[int, ...]:
        """Calculate the split size for each dimension based on the input shape and number of splits.

        Args:
            input_shape: Shape of the input tensor. If a tensor is provided, its shape will be used.

        Returns:
            Tuple of split sizes for each dimension.
        """
        input_shape = self.get_input_shape(input_shape)
        if num_splits is None:
            num_splits = self.num_splits

        context = self.get_edge_context()

        split_size = []
        for i in range(self.spatial_dims):
            split_size.append(input_shape[i] // num_splits[i] + 2 * context[i])
        split_size = tuple(split_size)
        return split_size

    def get_split_stride(
        self, input_shape: tuple[int, ...] | torch.Size | torch.Tensor, num_splits: tuple[int, ...] = None
    ) -> tuple[int, ...]:
        """Calculate the split stride for each dimension based on the input shape and context size.

        Args:
            input_shape: Shape of the input tensor. If a tensor is provided, its shape will be used.
            num_splits: Number of splits for each spatial dimension. If None, the default num_splits will be used.

        Returns:
            Tuple of split strides for each dimension.
        """
        input_shape = self.get_input_shape(input_shape)
        if num_splits is None:
            num_splits = self.num_splits
        context = self.get_edge_context()
        split_size = self.get_split_size(input_shape, num_splits)
        split_stride = [split_size[i] - 2 * context[i] for i in range(self.spatial_dims)]
        assert all(
            split_stride[i] > 0 for i in range(self.spatial_dims)
        ), "Split stride must be greater than 0 for all dimensions."
        return split_stride

    def pad_input_for_context(self, x: torch.Tensor) -> torch.Tensor:
        """Pad the input with the context size for consistent merging.

        Args:
            x: Input tensor of shape (batch_size, in_channels, [z], y, x).

        Returns:
            Padded input tensor of shape (batch_size, in_channels, [z], y, x).
        """
        context = self.get_edge_context()
        padding = [0, 0] * (x.ndim - self.spatial_dims)
        for i in range(self.spatial_dims):
            padding.extend([context[i], context[i]])
        x = F.pad(x, list(reversed(padding)))
        return x

    @populate_docstring
    def forward(self, x: torch.Tensor, channels_first: bool = True) -> torch.Tensor:
        """Forward pass through the convolution layer with tensor splitting parallelism. Main convolution occurs on it's
             device, but the output is built on the input tensor's device.

        Args:
            x: {INPUT_3D_DOC}

        Returns:
            {OUTPUT_3D_DOC}
        """
        input_device = x.device

        x = rearrange_channels(x, channels_first, True)
        B, DIMS = x.shape[0], x.shape[2:]  # (batch_size, in_channels, [z], y, x)

        # Optimize num_splits
        num_splits = self.num_splits
        if self.optimize_num_splits:
            num_splits = self.get_optimized_num_splits(x)

        # Pad input such that it is divisible by the number of splits
        x = self.pad_input_for_divisibility(x, num_splits)

        # Calculate split size
        split_size = self.get_split_size(x, num_splits)

        # Identify the stride required to split the input tensor such that overlapping regions can be counted only once
        split_stride = self.get_split_stride(x, num_splits)

        # Pad the input
        x = self.pad_input_for_context(x)

        # Split the input tensor
        splitter = Splitter(
            split_dims=self.spatial_dims,
            split_size=split_size,
            stride=split_stride,
            extend_mode=None,  # Padding has been handled here for better control
        )
        x = splitter(x)
        # (num_splits, batch_size, in_channels, [z1], y1, x1)

        # Run the convolution on each split
        outputs: list[torch.Tensor] = []
        for x_split, position in x:
            x_split = x_split.to(self.conv.weight.device)
            x_split = self.conv(x_split)
            x_split = x_split.to(input_device)
            outputs.append((x_split, position))
        # list of len num_splits, each of a tuple with the split of shape (batch_size, out_channels, [z1], y1, x1) and
        # the position as a tensor

        # Merge the outputs
        context = self.get_edge_context()
        merged = torch.empty((B, outputs[0][0].shape[1], *DIMS), device=input_device)
        for output, position in outputs:
            merged_slices = [slice(None), slice(None)]  # To track the coordinates where the output will be placed
            output_slices = [slice(None), slice(None)]  # To track the actual output that should be placed
            for i in range(self.spatial_dims):
                dim = i + 2
                merged_slice = slice(position[i], min(position[i] + split_stride[i], DIMS[i]))
                output_slice = slice(context[i], -context[i] if context[i] != 0 else None)

                merged_indices = merged_slice.indices(merged.shape[dim])
                output_indices = output_slice.indices(output.shape[dim])
                len_merged_slice = merged_indices[1] - merged_indices[0]
                len_output_slice = output_indices[1] - output_indices[0]
                if len_output_slice > len_merged_slice:
                    output_slice = slice(output_slice.start, output_slice.start + len_merged_slice)

                merged_slices.append(merged_slice)
                output_slices.append(output_slice)

            merged[tuple(merged_slices)] = output[tuple(output_slices)]

        merged = rearrange_channels(merged, True, channels_first)

        return merged

    def extra_repr(self):
        return f"num_splits={self.num_splits}"

# %% ../../nbs/blocks/04_cnn.ipynb 24
def add_tsp_to_module(
    module: nn.Module,
    num_splits_2d: int | tuple[int, int] | None = None,
    num_splits_3d: int | tuple[int, int, int] = None,
    strict: bool = True,
) -> nn.Module:
    """Recursively add TensorSplittingConv to the module for all Conv2d and Conv3d layers.

    Args:
        module: The module to modify.
        num_splits_2d: Number of splits for 2D convolutions. If None, 2D convolutions will not be modified.
        num_splits_3d: Number of splits for 3D convolutions. If None, 3D convolutions will not be modified.
        strict: Whether to raise an error if a conversion fails. If False, it will log the error and continue.

    Returns:
        The modified module with TensorSplittingConv layers.

    Raises:
        ValueError: If both num_splits_2d and num_splits_3d are None.
        Exception: If a conversion fails and strict is True.
    """

    if num_splits_2d is None and num_splits_3d is None:
        raise ValueError("At least one of num_splits_2d or num_splits_3d must be provided.")
    for name, child in module.named_children():
        if isinstance(child, TensorSplittingConv):
            continue
        if num_splits_2d is not None and isinstance(child, nn.Conv2d):
            try:
                setattr(module, name, TensorSplittingConv(child, num_splits_2d).to(child.weight.device))
            except Exception as e:
                if strict:
                    raise e
                else:
                    logger.debug(f"Could not convert {name} to TensorSplittingConv. Error: {e}")

        if num_splits_3d is not None and isinstance(child, nn.Conv3d):
            try:
                setattr(module, name, TensorSplittingConv(child, num_splits_3d).to(child.weight.device))
            except Exception as e:
                if strict:
                    raise e
                else:
                    logger.debug(f"Could not convert {name} to TensorSplittingConv. Error: {e}")
        else:
            add_tsp_to_module(child, num_splits_2d, num_splits_3d, strict)
    return module

# %% ../../nbs/blocks/04_cnn.ipynb 25
def remove_tsp_from_module(module: nn.Module) -> nn.Module:
    """Recursively remove TensorSplittingConv from the module and replace it with the original convolution layer.

    Args:
        module: The module to modify.

    Returns:
        The modified module with TensorSplittingConv layers replaced by the original convolution layers.
    """
    for name, child in module.named_children():
        if isinstance(child, TensorSplittingConv):
            setattr(module, name, child.conv.to(child.conv.weight.device))
        else:
            remove_tsp_from_module(child)
    return module
