# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/nets/01_swin_3d.ipynb.

# %% auto 0
__all__ = ['Swin3DPatchMergingConfig', 'Swin3DPatchSplittingConfig', 'Swin3DBlockConfig', 'Swin3DStageConfig',
           'Swin3DEncoderDecoderConfig', 'Swin3DEncoderWithPatchEmbeddingsConfig', 'Swin3DLayer', 'Swin3DBlock',
           'Swin3DPatchMerging', 'Swin3DPatchSplitting', 'Swin3DStage', 'Swin3DEncoderDecoderBase', 'Swin3DEncoder',
           'Swin3DDecoder', 'Swin3DEncoderWithPatchEmbeddings']

# %% ../../nbs/nets/01_swin_3d.ipynb 2
from functools import wraps

import numpy as np
import torch
from einops import rearrange
from huggingface_hub import PyTorchModelHubMixin
from loguru import logger
from torch import nn

from ..blocks.transformer import Attention3DWithMLP, Attention3DWithMLPConfig
from ..docstrings import populate_docstring
from vision_architectures.layers.embeddings import (
    AbsolutePositionEmbeddings3D,
    PatchEmbeddings3D,
    RelativePositionEmbeddings3D,
    RelativePositionEmbeddings3DConfig,
)
from ..utils.activation_checkpointing import ActivationCheckpointing
from ..utils.custom_base_model import CustomBaseModel, Field, computed_field, model_validator
from ..utils.rearrange import rearrange_channels

# %% ../../nbs/nets/01_swin_3d.ipynb 4
class Swin3DPatchMergingConfig(CustomBaseModel):
    in_dim: int = Field(..., description="Input dimension before merging")
    out_dim: int = Field(..., description="Output dimension after merging")
    merge_window_size: tuple[int, int, int] = Field(..., description="Size of the window for merging patches")

    @computed_field(description="Factor by which the dimension is increased after merging")
    @property
    def out_dim_ratio(self) -> float:
        return self.out_dim / self.in_dim

    @model_validator(mode="before")
    @classmethod
    def validate_before(cls, data):
        super().validate_before(data)
        merge_window_size = data.get("merge_window_size")
        if isinstance(merge_window_size, int):
            data["merge_window_size"] = (
                merge_window_size,
                merge_window_size,
                merge_window_size,
            )
        return data


class Swin3DPatchSplittingConfig(CustomBaseModel):
    in_dim: int = Field(..., description="Input dimension before splitting")
    out_dim: int = Field(..., description="Output dimension after splitting")
    final_window_size: tuple[int, int, int] = Field(..., description="Size of the window to split patches into")

    @computed_field(description="Factor by which the dimension is decreased after splitting")
    @property
    def in_dim_ratio(self) -> float:
        return self.in_dim / self.out_dim

    @model_validator(mode="before")
    @classmethod
    def validate_before(cls, data):
        super().validate_before(data)
        final_window_size = data.get("final_window_size")
        if isinstance(final_window_size, int):
            data["final_window_size"] = (
                final_window_size,
                final_window_size,
                final_window_size,
            )
        return data


class Swin3DBlockConfig(Attention3DWithMLPConfig):
    window_size: tuple[int, int, int] = Field(..., description="Size of the window to apply attention over")

    use_relative_position_bias: bool = Field(False, description="Whether to use relative position bias")
    patch_merging: Swin3DPatchMergingConfig | None = Field(
        None, description="Patch merging config if desired. Patch merging is applied before attention."
    )
    patch_splitting: Swin3DPatchSplittingConfig | None = Field(
        None, description="Patch splitting config if desired. Patch splitting is applied after attention."
    )

    in_dim: int | None = Field(None, description="Input dimension of the stage. Useful if ``patch_merging`` is used.")
    dim: int = Field(..., description="Dim at which attention is performed")
    out_dim: int | None = Field(
        None, description="Output dimension of the stage. Useful if ``patch_splitting`` is used."
    )

    @property
    def spatial_compression_ratio(self):
        compression_ratio = (1.0, 1.0, 1.0)
        if self.patch_merging is not None:
            compression_ratio = tuple(compression_ratio[i] * self.patch_merging.merge_window_size[i] for i in range(3))
        if self.patch_splitting is not None:
            compression_ratio = tuple(
                compression_ratio[i] / self.patch_splitting.final_window_size[i] for i in range(3)
            )
        return compression_ratio

    def get_out_patch_size(self, in_patch_size: tuple[int, int, int]):
        patch_size = tuple(int(in_patch_size[i] * self.spatial_compression_ratio[i]) for i in range(3))
        return patch_size

    def get_in_patch_size(self, out_patch_size: tuple[int, int, int]):
        patch_size = tuple(int(out_patch_size[i] / self.spatial_compression_ratio[i]) for i in range(3))
        return patch_size

    def get_in_dim(self) -> int:
        if self.in_dim is None:
            return self.dim
        return self.in_dim

    def get_out_dim(self) -> int:
        if self.out_dim is None:
            return self.dim
        return self.out_dim

    @property
    def out_dim_ratio(self) -> float:
        return self.get_out_dim() / self.get_in_dim()

    def populate(self):
        """Populate the in_dim and out_dim of patch_splitting and patch_merging based on the stage's in_dim, dim,
        out_dim."""
        if self.patch_merging is not None:
            if self.in_dim != self.patch_merging.in_dim:
                if self.in_dim is not None:
                    logger.warning(
                        f"Overwriting in_dim ({self.in_dim}) for this stage as it does not match patch_merging config "
                        f"in_dim ({self.patch_merging.in_dim})."
                    )
                self.in_dim = self.patch_merging.in_dim

            if self.dim != self.patch_merging.out_dim:
                if self.dim is not None:
                    logger.warning(
                        f"Overwriting dim ({self.dim}) for this stage as it does not match patch_merging config out_dim "
                        f"({self.patch_merging.out_dim})."
                    )
                self.dim = self.patch_merging.out_dim

        if self.patch_splitting is not None:
            if self.dim != self.patch_splitting.in_dim:
                if self.dim is not None:
                    logger.warning(
                        f"Overwriting dim ({self.dim}) for this stage as it does not match patch_splitting config "
                        f"in_dim ({self.patch_splitting.in_dim})."
                    )
                self.dim = self.patch_splitting.in_dim

            if self.out_dim != self.patch_splitting.out_dim:
                if self.out_dim is not None:
                    logger.warning(
                        f"Overwriting out_dim ({self.out_dim}) for this stage as it does not match patch_splitting config "
                        f"out_dim ({self.patch_splitting.out_dim})."
                    )
                self.out_dim = self.patch_splitting.out_dim

    @model_validator(mode="after")
    def validate(self):
        super().validate()
        self.populate()
        return self


class Swin3DStageConfig(Swin3DBlockConfig):
    depth: int = Field(..., description="Number of transformer blocks in this stage")


class Swin3DEncoderDecoderConfig(CustomBaseModel):
    stages: list[Swin3DStageConfig]

    def populate(self):
        """Populate the in_dim, dim, out_dim of each stage."""
        for i in range(len(self.stages)):
            self.stages[i].populate()

    @model_validator(mode="after")
    def validate(self):
        super().validate()
        self.populate()

        # Ensure there is at least one stage
        assert len(self.stages) > 0, "Must have at least one stage"

        # Test divisibility of dim with number of attention heads
        for stage in self.stages:
            assert (
                stage.dim % stage.num_heads == 0
            ), f"stage.dim {stage.dim} is not divisible by stage.num_heads {stage.num_heads}"

        # Ensure dimensionality matches across stages
        for i in range(len(self.stages) - 1):
            prev_out_dim = self.stages[i].get_out_dim()
            succ_in_dim = self.stages[i + 1].get_in_dim()
            assert prev_out_dim == succ_in_dim, (
                f"Dimensionality mismatch between stages. Preceding stage has out_dim "
                f"{prev_out_dim} and succeeding stage has in_dim {succ_in_dim}."
            )

        return self

    def get_out_dim_ratios(self):
        return [stage.out_dim_ratio for stage in self.stages]


class Swin3DEncoderWithPatchEmbeddingsConfig(Swin3DEncoderDecoderConfig):
    in_channels: int = Field(..., description="Number of input channels in the input image/video")
    patch_size: tuple[int, int, int] = Field(
        ..., description="Size of the patches to be extracted from the input image/video"
    )
    image_size: tuple[int, int, int] | None = Field(
        None, description="Size of the input image/video. Required if absolute position embeddings are learnable."
    )

    use_absolute_position_embeddings: bool = Field(True, description="Whether to use absolute position embeddings.")
    learnable_absolute_position_embeddings: bool = Field(
        False, description="Whether to use learnable absolute position embeddings."
    )

    @model_validator(mode="after")
    def validate(self):
        super().validate()
        # Test population of image_size field iff the absolute position embeddings are relative
        if self.learnable_absolute_position_embeddings:
            assert (
                self.image_size is not None
            ), "Please provide image_size if absolute position embeddings are learnable"
        return self

# %% ../../nbs/nets/01_swin_3d.ipynb 8
@populate_docstring
class Swin3DLayer(nn.Module):
    """Swin 3D Layer applying windowed attention with optional relative position embeddings.
    {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(
        self,
        config: RelativePositionEmbeddings3DConfig | Attention3DWithMLPConfig = {},
        checkpointing_level: int = 0,
        **kwargs
    ):
        """Initializes the Swin3DLayer.

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

        if isinstance(config, CustomBaseModel):
            config = config.model_dump()
        self._all_kwargs = config | kwargs
        self._window_size = self._all_kwargs.get("window_size")
        self._use_relative_position_bias = self._all_kwargs.get("use_relative_position_bias")

        self.embeddings_config = RelativePositionEmbeddings3DConfig.model_validate(
            self._all_kwargs | {"grid_size": self._window_size}
        )
        self.transformer_config = Attention3DWithMLPConfig.model_validate(self._all_kwargs)

        relative_position_bias = None
        if self._use_relative_position_bias:
            relative_position_bias = RelativePositionEmbeddings3D(self.embeddings_config)

        self.transformer = Attention3DWithMLP(
            self.transformer_config,
            relative_position_bias=relative_position_bias,
            checkpointing_level=checkpointing_level,
        )

        self.checkpointing_level3 = ActivationCheckpointing(3, checkpointing_level)

    @staticmethod
    def _get_rearrange_patterns() -> tuple[str, str]:
        """Note that the patterns will be applied on tensors that are in channels_last format"""
        forward_pattern = (
            "b (num_windows_z window_size_z) (num_windows_y window_size_y) (num_windows_x window_size_x) dim -> "
            "(b num_windows_z num_windows_y num_windows_x) window_size_z window_size_y window_size_x dim "
        )
        reverse_pattern = (
            "(b num_windows_z num_windows_y num_windows_x) window_size_z window_size_y window_size_x dim -> "
            "b (num_windows_z window_size_z) (num_windows_y window_size_y) (num_windows_x window_size_x) dim"
        )
        return forward_pattern, reverse_pattern

    def _forward(self, hidden_states: torch.Tensor, channels_first: bool = True) -> torch.Tensor:
        """Window the input features and apply self attention on each window.

        Args:
            hidden_states: {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            {OUTPUT_3D_DOC}
        """
        # hidden_states: (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        hidden_states = rearrange_channels(hidden_states, channels_first, False)
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        _, num_patches_z, num_patches_y, num_patches_x, _ = hidden_states.shape

        forward_pattern, reverse_pattern = self._get_rearrange_patterns()

        # Perform windowing
        window_size_z, window_size_y, window_size_x = self._window_size
        num_windows_z, num_windows_y, num_windows_x = (
            num_patches_z // window_size_z,
            num_patches_y // window_size_y,
            num_patches_x // window_size_x,
        )
        hidden_states = rearrange(
            hidden_states,
            forward_pattern,
            num_windows_z=num_windows_z,
            num_windows_y=num_windows_y,
            num_windows_x=num_windows_x,
            window_size_z=window_size_z,
            window_size_y=window_size_y,
            window_size_x=window_size_x,
        ).contiguous()

        hidden_states = self.transformer(hidden_states, hidden_states, hidden_states, channels_first=False)

        # Undo windowing
        output = rearrange(
            hidden_states,
            reverse_pattern,
            num_windows_z=num_windows_z,
            num_windows_y=num_windows_y,
            num_windows_x=num_windows_x,
            window_size_z=window_size_z,
            window_size_y=window_size_y,
            window_size_x=window_size_x,
        ).contiguous()

        output = rearrange_channels(output, False, channels_first)
        # (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        return output

    @wraps(_forward)
    def forward(self, *args, **kwargs) -> torch.Tensor:
        return self.checkpointing_level3(self._forward, *args, **kwargs)

# %% ../../nbs/nets/01_swin_3d.ipynb 11
@populate_docstring
class Swin3DBlock(nn.Module):
    """Swin 3D Block consisting of two Swin3DLayers: one with regular windows and one with shifted windows.
    {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: Swin3DBlockConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initializes the Swin3DBlock.

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

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

        self.w_layer = Swin3DLayer(self.config.model_dump(), checkpointing_level=checkpointing_level)
        self.sw_layer = Swin3DLayer(self.config.model_dump(), checkpointing_level=checkpointing_level)

    @populate_docstring
    def forward(
        self, hidden_states: torch.Tensor, channels_first: bool = True, return_intermediates: bool = False
    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
        """Apply window attention and shifted window attention on the input features.

        Args:
            hidden_states: {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}
            return_intermediates: {RETURN_INTERMEDIATES_DOC}

        Returns:
            {OUTPUT_3D_DOC}. If return_intermediates is True, also returns a list of intermediate layer outputs. Note
            that the intermediate layer outputs returned will always be in ``channels_last`` format.
        """
        # hidden_states: (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        hidden_states = rearrange_channels(hidden_states, channels_first, False)
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        layer_outputs = []

        # First layer
        hidden_states = self.w_layer(hidden_states, channels_first=False)
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        layer_outputs.append(hidden_states)

        # Shift windows
        window_size_z, window_size_y, window_size_x = self.config.window_size
        shifts = (window_size_z // 2, window_size_y // 2, window_size_x // 2)
        hidden_states = torch.roll(hidden_states, shifts=shifts, dims=(1, 2, 3))
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        # Second layer
        hidden_states = self.sw_layer(hidden_states, channels_first=False)
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        # Reverse window shift
        shifts = tuple(-shift for shift in shifts)
        hidden_states = torch.roll(hidden_states, shifts=shifts, dims=(1, 2, 3))
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        layer_outputs.append(hidden_states)

        hidden_states = rearrange_channels(hidden_states, False, channels_first)
        # (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        if return_intermediates:
            return hidden_states, layer_outputs
        return hidden_states

# %% ../../nbs/nets/01_swin_3d.ipynb 13
@populate_docstring
class Swin3DPatchMerging(nn.Module):
    """Patch merging layer for Swin3D. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: Swin3DPatchMergingConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the Swin3DPatchMerging layer.

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

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

        in_dim = self.config.in_dim * np.prod(self.config.merge_window_size)
        self.layer_norm = nn.LayerNorm(in_dim)
        self.proj = nn.Linear(in_dim, self.config.out_dim)

        self.checkpointing_level1 = ActivationCheckpointing(1, checkpointing_level)

    @populate_docstring
    def _forward(self, hidden_states: torch.Tensor, channels_first: bool = True) -> torch.Tensor:
        """Merge multiple patches into a single patch.

        Args:
            hidden_states: {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            {OUTPUT_3D_DOC}
        """
        # hidden_states: (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        hidden_states = rearrange_channels(hidden_states, channels_first, False)
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        window_size_z, window_size_y, window_size_x = self.config.merge_window_size

        hidden_states = rearrange(
            hidden_states,
            "b (new_num_patches_z window_size_z) (new_num_patches_y window_size_y) (new_num_patches_x window_size_x) dim -> "
            "b new_num_patches_z new_num_patches_y new_num_patches_x (window_size_z window_size_y window_size_x dim)",
            window_size_z=window_size_z,
            window_size_y=window_size_y,
            window_size_x=window_size_x,
        ).contiguous()

        hidden_states = self.layer_norm(hidden_states)
        hidden_states = self.proj(hidden_states)

        hidden_states = rearrange_channels(hidden_states, False, channels_first)
        # (b, [dim], new_num_patches_z, new_num_patches_y, new_num_patches_x, [dim])

        return hidden_states

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

# %% ../../nbs/nets/01_swin_3d.ipynb 15
@populate_docstring
class Swin3DPatchSplitting(nn.Module):
    """Patch splitting layer for Swin3D. {CLASS_DESCRIPTION_3D_DOC}

    This is a self-implemented class and is not part of the paper."""

    @populate_docstring
    def __init__(self, config: Swin3DPatchSplittingConfig, checkpointing_level: int = 0, **kwargs):
        """Initialize the Swin3DPatchSplitting layer.

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

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

        self.layer_norm = nn.LayerNorm(self.config.in_dim)
        self.proj = nn.Linear(self.config.in_dim, self.config.out_dim * np.prod(self.config.final_window_size))

        self.checkpointing_level1 = ActivationCheckpointing(1, checkpointing_level)

    @populate_docstring
    def _forward(self, hidden_states: torch.Tensor, channels_first: bool = True) -> torch.Tensor:
        """Split patches into multiple patches.

        Args:
            hidden_states: {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}

        Returns:
            {OUTPUT_3D_DOC}
        """
        # hidden_states: (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        hidden_states = rearrange_channels(hidden_states, channels_first, False)
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        hidden_states = self.layer_norm(hidden_states)
        hidden_states = self.proj(hidden_states)

        window_size_z, window_size_y, window_size_x = self.config.final_window_size

        hidden_states = rearrange(
            hidden_states,
            "b num_patches_z num_patches_y num_patches_x (window_size_z window_size_y window_size_x dim) -> "
            "b (num_patches_z window_size_z) (num_patches_y window_size_y) (num_patches_x window_size_x) dim",
            window_size_z=window_size_z,
            window_size_y=window_size_y,
            window_size_x=window_size_x,
        ).contiguous()

        hidden_states = rearrange_channels(hidden_states, False, channels_first)
        # (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        return hidden_states

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

# %% ../../nbs/nets/01_swin_3d.ipynb 17
@populate_docstring
class Swin3DStage(nn.Module):
    """Swin3D stage for Swin3D. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: Swin3DStageConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initialize the Swin3DStage.

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

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

        self.patch_merging = None
        if self.config.patch_merging is not None:
            self.patch_merging = Swin3DPatchMerging(self.config.patch_merging)

        self.blocks = nn.ModuleList(
            [Swin3DBlock(self.config) for _ in range(self.config.depth)],
        )

        self.patch_splitting = None
        if self.config.patch_splitting is not None:
            # This has been implemented to create a Swin-based decoder
            self.patch_splitting = Swin3DPatchSplitting(self.config.patch_splitting)

        self.checkpointing_level4 = ActivationCheckpointing(4, checkpointing_level)

    @populate_docstring
    def _forward(
        self, hidden_states: torch.Tensor, channels_first: bool = True, return_intermediates: bool = False
    ) -> torch.Tensor:
        """Merge patches if applicable (used by the encoder), perform a series of window and shifted window attention,
        and then split patches if applicable (used by the decoder).

        Args:
            hidden_states: {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}
            return_intermediates: {RETURN_INTERMEDIATES_DOC}

        Returns:
            {OUTPUT_3D_DOC}. If return_intermediates is True, also returns a list of intermediate layer outputs. Note
            that the intermediate layer outputs returned will always be in ``channels_last`` format.
        """
        # hidden_states: (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        hidden_states = rearrange_channels(hidden_states, channels_first, False)
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        if self.patch_merging:
            hidden_states = self.patch_merging(hidden_states, channels_first=False)
            # (b, new_num_patches_z, new_num_patches_y, new_num_patches_x, new_dim)

        layer_outputs = []
        for layer_module in self.blocks:
            hidden_states, _layer_outputs = layer_module(hidden_states, channels_first=False, return_intermediates=True)
            # (b, new_num_patches_z, new_num_patches_y, new_num_patches_x, new_dim)
            layer_outputs.extend(_layer_outputs)

        if self.patch_splitting:
            hidden_states = self.patch_splitting(hidden_states, channels_first=False)
            # (b, new_num_patches_z, new_num_patches_y, new_num_patches_x, new_dim)

        hidden_states = rearrange_channels(hidden_states, False, channels_first)
        # (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        if return_intermediates:
            return hidden_states, layer_outputs
        return hidden_states

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

# %% ../../nbs/nets/01_swin_3d.ipynb 20
class Swin3DEncoderDecoderBase(nn.Module, PyTorchModelHubMixin):
    @populate_docstring
    def __init__(self, config: Swin3DEncoderDecoderConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initializes the Swin3DEncoder/Swin3DDecoder.

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

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

        self.stages = nn.ModuleList(
            [Swin3DStage(stage_config, checkpointing_level) for stage_config in self.config.stages]
        )

        self.checkpointing_level5 = ActivationCheckpointing(5, checkpointing_level)

    @populate_docstring
    def _forward(
        self, hidden_states: torch.Tensor, channels_first: bool = True, return_intermediates: bool = False
    ) -> torch.Tensor:
        """Encodes the input features using the Swin Transformer hierarchy.

        Args:
            hidden_states: {INPUT_3D_DOC}
            channels_first: {CHANNELS_FIRST_DOC}
            return_intermediates: {RETURN_INTERMEDIATES_DOC}

        Returns:
            {OUTPUT_3D_DOC}. If return_intermediates is True, also returns a list of intermediate layer outputs. Note
            that the intermediate layer outputs returned will always be in ``channels_last`` format.
        """
        # hidden_states: (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        hidden_states = rearrange_channels(hidden_states, channels_first, False)
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        stage_outputs, layer_outputs = [], []
        for stage_module in self.stages:
            hidden_states, _layer_outputs = stage_module(hidden_states, channels_first=False, return_intermediates=True)
            # (b, new_num_patches_z, new_num_patches_y, new_num_patches_x, dim)

            stage_outputs.append(hidden_states)
            layer_outputs.extend(_layer_outputs)

        hidden_states = rearrange_channels(hidden_states, False, channels_first)
        # (b, [dim], num_patches_z, num_patches_y, num_patches_x, [dim])

        if return_intermediates:
            return hidden_states, stage_outputs, layer_outputs
        return hidden_states

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

# %% ../../nbs/nets/01_swin_3d.ipynb 22
@populate_docstring
class Swin3DEncoder(Swin3DEncoderDecoderBase):
    """3D Swin Transformer encoder. Assumes input has already been patchified/tokenized. {CLASS_DESCRIPTION_3D_DOC}"""

    def __init__(self, config: Swin3DEncoderDecoderConfig = {}, checkpointing_level: int = 0, **kwargs):
        super().__init__(config, checkpointing_level, **kwargs)

        for stage_config in self.config.stages:
            if stage_config.patch_splitting is not None:
                assert stage_config.patch_merging is not None, "Swin3DEncoder is not for decoding (mid blocks are ok)."

# %% ../../nbs/nets/01_swin_3d.ipynb 25
@populate_docstring
class Swin3DDecoder(Swin3DEncoderDecoderBase):
    """3D Swin Transformer decoder. Assumes input has already been patchified/tokenized. {CLASS_DESCRIPTION_3D_DOC}"""

    def __init__(self, config: Swin3DEncoderDecoderConfig = {}, checkpointing_level: int = 0, **kwargs):
        super().__init__(config, checkpointing_level, **kwargs)

        for stage_config in config.stages:
            if stage_config.patch_merging is not None:
                assert (
                    stage_config.patch_splitting is not None
                ), "Swin3DDecoder is not for encoding (mid blocks are ok)."

# %% ../../nbs/nets/01_swin_3d.ipynb 28
@populate_docstring
class Swin3DEncoderWithPatchEmbeddings(nn.Module, PyTorchModelHubMixin):
    """3D Swin transformer with 3D patch embeddings. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: Swin3DEncoderWithPatchEmbeddingsConfig = {}, checkpointing_level: int = 0, **kwargs):
        """Initializes the Swin3DEncoderWithPatchEmbeddings.

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

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

        self.patchify = PatchEmbeddings3D(
            patch_size=self.config.patch_size,
            in_channels=self.config.in_channels,
            dim=self.config.stages[0].get_in_dim(),
            checkpointing_level=checkpointing_level,
        )
        self.absolute_position_embeddings = AbsolutePositionEmbeddings3D(
            dim=self.config.stages[0].get_in_dim(), learnable=False
        )
        self.encoder = Swin3DEncoder(self.config, checkpointing_level=checkpointing_level)

    @populate_docstring
    def forward(
        self,
        pixel_values: torch.Tensor,
        spacings: torch.Tensor = None,
        crop_offsets: torch.Tensor = None,
        channels_first: bool = True,
        return_intermediates: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor], list[torch.Tensor]]:
        """Patchify the input pixel values and then pass it through the Swin transformer.

        Args:
            pixel_values: {INPUT_3D_DOC}
            spacings: {SPACINGS_DOC}
            crop_offsets: Used if the embeddings required are of a crop of a larger image. If provided, the grid
                coordinates will be offset accordingly.
            channels_first: {CHANNELS_FIRST_DOC}
            return_intermediates: {RETURN_INTERMEDIATES_DOC}

        Returns:
            {OUTPUT_3D_DOC}. If `return_intermediates` is True, also returns the intermediate stage outputs and layer
            outputs.
        """
        # pixel_values: (b, [c], z, y, x, [c])
        # spacings: (b, 3)

        pixel_values = rearrange_channels(pixel_values, channels_first, False)
        # (b, z, y, x, c)

        embeddings = self.patchify(pixel_values, channels_first=False)
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        embeddings = self.absolute_position_embeddings(
            embeddings, spacings=spacings, crop_offsets=crop_offsets, channels_first=False
        )
        # (b, num_patches_z, num_patches_y, num_patches_x, dim)

        encoded, stage_outputs, layer_outputs = self.encoder(
            embeddings, channels_first=False, return_intermediates=True
        )
        # encoded: (b, new_num_patches_z, new_num_patches_y, new_num_patches_x, dim)
        # stage_outputs, layer_outputs: list of (b, some_num_patches_z, some_num_patches_y, some_num_patches_x, dim)

        encoded = rearrange_channels(encoded, False, channels_first)
        # (b [dim], new_num_patches_z, new_num_patches_y, new_num_patches_x, [dim])

        if return_intermediates:
            return encoded, stage_outputs, layer_outputs
        return encoded
