# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/layers/02_embeddings.ipynb.

# %% auto 0
__all__ = ['RelativePositionEmbeddings', 'get_absolute_position_embeddings_3d', 'get_timestep_embeddings_1d',
           'get_all_timestep_embeddings_1d', 'get_absolute_position_embeddings_1d',
           'RelativePositionEmbeddings3DConfig', 'AbsolutePositionEmbeddings3DConfig',
           'AbsolutePositionEmbeddings1DConfig', 'PatchEmbeddings3DConfig', 'get_coords_grid',
           'RelativePositionEmbeddings3D', 'RelativePositionEmbeddings3DMetaNetwork', 'get_sinusoidal_embeddings_3d',
           'AbsolutePositionEmbeddings3D', 'get_specific_sinusoidal_embeddings_1d', 'get_sinusoidal_embeddings_1d',
           'AbsolutePositionEmbeddings1D', 'PatchEmbeddings3D']

# %% ../../nbs/layers/02_embeddings.ipynb 2
from typing import Literal, Union

import numpy as np
import torch
from einops import rearrange, repeat
from torch import nn

from ..blocks.cnn import CNNBlock3D, CNNBlockConfig
from ..docstrings import populate_docstring
from ..utils.activation_checkpointing import ActivationCheckpointing
from ..utils.custom_base_model import CustomBaseModel, Field, model_validator
from ..utils.rearrange import rearrange_channels

# %% ../../nbs/layers/02_embeddings.ipynb 4
class RelativePositionEmbeddings3DConfig(CustomBaseModel):
    num_heads: int = Field(..., description="Number of query attention heads")
    grid_size: tuple[int, int, int] = Field(..., description="Size of entire patch matrix.")

    @property
    def num_patches(self) -> int:
        """Number of patches."""
        return np.prod(self.grid_size)

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

    @model_validator(mode="after")
    def validate(self):
        super().validate()
        if isinstance(self.grid_size, int):
            self.grid_size = (self.grid_size, self.grid_size, self.grid_size)
        return self


class AbsolutePositionEmbeddings3DConfig(CustomBaseModel):
    dim: int | None = Field(None, description="Dimension of the position embeddings")
    grid_size: tuple[int, int, int] | None = Field(None, description="Size of entire patch matrix.")
    learnable: bool = Field(False, description="Whether the position embeddings are learnable.")

    @property
    def num_patches(self) -> int:
        """Number of patches."""
        if self.grid_size is None:
            return None
        return np.prod(self.grid_size)

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

    @model_validator(mode="after")
    def validate(self):
        super().validate()
        if self.learnable and (self.dim is None or self.grid_size is None):
            raise ValueError("dim and grid_size must be provided if learnable is True")
        return self


class AbsolutePositionEmbeddings1DConfig(CustomBaseModel):
    dim: int | None = Field(None, description="Dimension of the position embeddings")
    length: int | None = Field(None, description="Length of the sequence.")
    learnable: bool = Field(False, description="Whether the position embeddings are learnable.")

    @model_validator(mode="after")
    def validate(self):
        super().validate()
        if self.learnable and (self.dim is None or self.length is None):
            raise ValueError("dim and length must be provided if learnable is True")
        return self


class PatchEmbeddings3DConfig(CNNBlockConfig):
    patch_size: tuple[int, int, int] = Field(..., description="Size of the patches to extract from the input.")
    in_channels: int = Field(..., description="Number of input channels.")
    dim: int = Field(..., description="Dimension of the embeddings.")
    norm_layer: str = Field("layernorm", description="Normalization layer to use.")

    out_channels: None = None
    kernel_size: None = None

    @model_validator(mode="before")
    @classmethod
    def validate_before(cls, data: dict):
        data.setdefault("patch_size", data.pop("kernel_size", None))
        data.setdefault("dim", data.pop("out_channels", None))
        return data

# %% ../../nbs/layers/02_embeddings.ipynb 7
def get_coords_grid(grid_size: tuple[int, int, int]) -> torch.Tensor:
    """Get a coordinate grid of shape (3, d, h, w) for a given grid size.

    Args:
        grid_size: Size of the grid (d, h, w).

    Returns:
        A tensor of shape (3, d, h, w) containing the coordinates.
    """
    d, h, w = grid_size

    grid_d = torch.arange(d, dtype=torch.int32)
    grid_h = torch.arange(h, dtype=torch.int32)
    grid_w = torch.arange(w, dtype=torch.int32)

    grid = torch.meshgrid(grid_d, grid_h, grid_w, indexing="ij")
    grid = torch.stack(grid, axis=0)
    # (3, d, h, w)

    return grid

# %% ../../nbs/layers/02_embeddings.ipynb 8
@populate_docstring
class RelativePositionEmbeddings3D(nn.Module):
    """Learnable 3D Relative Position Embeddings. This can be passed directly to the attention layers.
    {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: RelativePositionEmbeddings3DConfig = {}, **kwargs):
        """Initialize RelativePositionEmbeddings3D.

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

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

        num_heads = self.config.num_heads
        grid_size = self.config.grid_size

        # TODO: Add embed_spacing_info functionality

        relative_limits = (
            2 * grid_size[0] - 1,
            2 * grid_size[1] - 1,
            2 * grid_size[2] - 1,
        )

        self.relative_position_bias_table = nn.Parameter(torch.randn(num_heads, np.prod(relative_limits)))
        # (num_heads, num_patches_z * num_patches_y * num_patches_x)

        # Pair-wise relative position index for each token inside the window
        coords = get_coords_grid(grid_size)
        coords_flatten = rearrange(coords, "three d h w -> three (d h w)", three=3).contiguous()
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
        relative_coords[:, :, 0] += grid_size[0] - 1
        relative_coords[:, :, 1] += grid_size[1] - 1
        relative_coords[:, :, 2] += grid_size[2] - 1
        relative_position_index: torch.Tensor = (
            relative_coords[:, :, 0] * relative_limits[1] * relative_limits[2]
            + relative_coords[:, :, 1] * relative_limits[2]
            + relative_coords[:, :, 2]
        )
        self.relative_position_index = relative_position_index.flatten()
        # (num_patches, num_patches)

    def forward(self) -> torch.Tensor:
        """Get relative position embeddings as specified by the config.

        Returns:
            A tensor of shape (1, num_heads, num_patches, num_patches) containing the relative position embeddings.
        """
        relative_position_embeddings = self.relative_position_bias_table[:, self.relative_position_index].contiguous()
        # (num_heads, num_patches, num_patches)
        relative_position_embeddings = relative_position_embeddings.reshape(
            1, self.config.num_patches, self.config.num_patches, -1
        )
        # (1, num_patches, num_patches, num_heads)
        relative_position_embeddings = rearrange(
            relative_position_embeddings,
            "1 num_patches1 num_patches2 num_heads -> 1 num_heads num_patches1 num_patches2",
        ).contiguous()
        # (1, num_heads, num_patches, num_patches)
        return relative_position_embeddings

# %% ../../nbs/layers/02_embeddings.ipynb 10
@populate_docstring
class RelativePositionEmbeddings3DMetaNetwork(nn.Module):
    """3D Relative Position Embeddings obtained from a meta network (inspired by SwinV2). This can be passed directly
    to the attention layers. {CLASS_DESCRIPTION_3D_DOC}"""

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

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

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

        num_heads = self.config.num_heads
        grid_size = self.config.grid_size

        # TODO: Add embed_spacing_info functionality
        self.cpb_mlp = nn.Sequential(
            nn.Linear(3, 512, bias=True),
            nn.ReLU(inplace=True),
            nn.Linear(512, num_heads, bias=False),
        )

        relative_limits = (
            2 * grid_size[0] - 1,
            2 * grid_size[1] - 1,
            2 * grid_size[2] - 1,
        )

        # Relative coordinates table
        relative_coords_table = get_coords_grid(relative_limits).float()
        for i in range(3):
            relative_coords_table[i] = (relative_coords_table[i] - (grid_size[0] - 1)) / (
                grid_size[0] - 1 + 1e-5  # small value added to ensure there is no NaN when window size is 1
            )
        relative_coords_table = rearrange(
            relative_coords_table,
            "three num_patches_z num_patches_y num_patches_x -> num_patches_z num_patches_y num_patches_x three",
        ).contiguous()
        relative_coords_table *= 8  # Normalize to -8, 8
        relative_coords_table = (
            torch.sign(relative_coords_table) * torch.log2(torch.abs(relative_coords_table) + 1.0) / np.log2(8)
        )
        # (num_patches_z, num_patches_y, num_patches_x, 3)
        # Allow moving this to and from cuda whenever required but don't save to state_dict
        self.register_buffer("relative_coords_table", relative_coords_table, persistent=False)

        # Pair-wise relative position index for each token inside the window
        coords = get_coords_grid(grid_size)
        coords_flatten = rearrange(coords, "three d h w -> three (d h w)", three=3).contiguous()
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
        relative_coords = rearrange(
            relative_coords, "three num_patches1 num_patches2 -> num_patches1 num_patches2 three"
        ).contiguous()
        relative_coords[:, :, 0] += grid_size[0] - 1
        relative_coords[:, :, 1] += grid_size[1] - 1
        relative_coords[:, :, 2] += grid_size[2] - 1
        relative_position_index: torch.Tensor = (
            relative_coords[:, :, 0] * relative_limits[1] * relative_limits[2]
            + relative_coords[:, :, 1] * relative_limits[2]
            + relative_coords[:, :, 2]
        )
        self.relative_position_index = relative_position_index.flatten()
        # (num_patches, num_patches)

        self.checkpointing_level1 = ActivationCheckpointing(1, checkpointing_level)

    def get_relative_position_embeddings_table(self) -> torch.Tensor:
        """Get the relative position embeddings table from the meta network.

        Returns:
            A tensor of shape (num_patches, num_heads) containing the relative position embeddings table.
        """
        # (num_patches_z, num_patches_y, num_patches_x, 3)
        relative_position_embeddings_table: torch.Tensor = self.cpb_mlp(self.relative_coords_table)
        # (num_patches_z, num_patches_y, num_patches_x, num_heads)
        relative_position_embeddings_table = relative_position_embeddings_table.reshape(-1, self.config.num_heads)
        # (num_patches, num_heads)
        return relative_position_embeddings_table

    def forward(self) -> torch.Tensor:
        """Get relative position embeddings as specified by the config.

        Returns:
            A tensor of shape (num_heads, num_patches, num_patches) containing the relative position embeddings.
        """
        relative_position_embeddings_table = self.checkpointing_level1(self.get_relative_position_embeddings_table)
        # (num_patches, num_heads)
        relative_position_embeddings = relative_position_embeddings_table[self.relative_position_index]
        # (num_patches * num_patches, num_heads)
        relative_position_embeddings = rearrange(
            relative_position_embeddings,
            "(num_patches1 num_patches2) num_heads -> num_heads num_patches1 num_patches2",
            num_patches1=self.config.num_patches,
            num_patches2=self.config.num_patches,
            num_heads=self.config.num_heads,
        ).contiguous()
        # (num_heads, num_patches, num_patches)
        relative_position_embeddings = 16 * torch.sigmoid(relative_position_embeddings)
        # (num_heads, num_patches, num_patches)
        return relative_position_embeddings

# %% ../../nbs/layers/02_embeddings.ipynb 12
RelativePositionEmbeddings = Union[RelativePositionEmbeddings3D, RelativePositionEmbeddings3DMetaNetwork]

# %% ../../nbs/layers/02_embeddings.ipynb 13
@populate_docstring
def get_sinusoidal_embeddings_3d(
    dim: int,
    grid_size: tuple[int, int, int],
    spacing: tuple[float, float, float] = (1.0, 1.0, 1.0),
    crop_offset: tuple[int, int, int] = None,
    channels_first: bool = True,
) -> torch.Tensor:
    """Get 3D sinusoidal position embeddings.

    Args:
        dim: Embedding dimension. Must be divisible by 6.
        grid_size: Size of the patch grid (d, h, w).
        spacing: Spacing between patches in each dimension. Useful for medical images.
        crop_offset: 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}

    Returns:
        {OUTPUT_3D_DOC}
    """
    if dim % 6 != 0:
        raise ValueError("dim must be divisible by 6")

    grid = get_coords_grid(grid_size)
    # (3, d, h, w)

    # Apply offset if crop parameters are provided
    if crop_offset is not None:
        # Offset the grid coordinates to represent their position in the full volume
        for i in range(3):
            grid[i] = grid[i] + crop_offset[i]

    grid = rearrange(grid, "x d h w -> x 1 d h w").contiguous()
    # (3, 1, d, h, w)

    omega = torch.arange(dim // 6, dtype=torch.float32)
    omega /= dim / 6.0
    omega = 1.0 / (10000**omega)
    # (dim // 6)

    patch_multiplier = torch.Tensor(spacing) / min(spacing)

    embeddings = []
    for i, grid_subset in enumerate(grid):
        grid_subset = grid_subset.reshape(-1)

        out = torch.einsum("m,d->md", grid_subset, omega)

        emb_sin = torch.sin(out)
        emb_cos = torch.cos(out)

        emb = torch.cat([emb_sin, emb_cos], axis=1) * patch_multiplier[i]
        embeddings.append(emb)

    embeddings = torch.cat(embeddings, axis=1)
    # (dim, d * h * w)
    embeddings = rearrange(
        embeddings,
        "(d h w) e -> 1 e d h w",
        d=grid_size[0],
        h=grid_size[1],
        w=grid_size[2],
    ).contiguous()
    # (1, dim, d, h, w)

    embeddings = rearrange_channels(embeddings, True, channels_first)
    # (1, [dim], d, h, w, [dim])

    return embeddings


get_absolute_position_embeddings_3d = get_sinusoidal_embeddings_3d

# %% ../../nbs/layers/02_embeddings.ipynb 14
@populate_docstring
class AbsolutePositionEmbeddings3D(nn.Module):
    """3D Absolute Position Embeddings. May or may not learnable. These have to be applied on the input manually and
    cannot be passed to attention layers directly. {CLASS_DESCRIPTION_3D_DOC}"""

    @populate_docstring
    def __init__(self, config: AbsolutePositionEmbeddings3DConfig = {}, **kwargs):
        """Initialize AbsolutePositionEmbeddings3D.

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

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

        dim = self.config.dim
        grid_size = self.config.grid_size
        learnable = self.config.learnable

        self.position_embeddings_cache = {}
        self.position_embeddings = None
        if dim is not None and grid_size is not None:
            self.position_embeddings = nn.Parameter(
                get_absolute_position_embeddings_3d(dim, grid_size), requires_grad=learnable
            )

    @populate_docstring
    def forward(
        self,
        x: torch.Tensor,
        embedding_type: Literal["add", "concat"] = "add",
        spacings: torch.Tensor = None,
        channels_first: bool = True,
        crop_offsets: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Apply absolute position embeddings to the input tensor.

        Args:
            x: {INPUT_3D_DOC}
            embedding_type: Type of embedding to apply. 'add' to add the position embeddings to the input,
                'concat' to concatenate them along the channel dimension.
            spacings: {SPACINGS_DOC}
            channels_first: {CHANNELS_FIRST_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.

        Returns:
            {OUTPUT_3D_DOC}
        """
        assert x.ndim == 5, "Input tensor must be of shape (b, [d], z, y, x, [d])"
        # Check if sufficient information has been provided
        if self.position_embeddings is None:
            dim = self.config.dim
            if dim is None:
                dim = x.shape[1] if channels_first else x.shape[-1]
            grid_size = self.config.grid_size
            if grid_size is None:
                grid_size = tuple(x.shape[2:5] if channels_first else x.shape[1:4])
        else:
            dim = self.config.dim
            grid_size = self.config.grid_size

        # Estimate batch size
        b = x.shape[0]

        # Get position embeddings, adjust based on crop offsets if applicable
        if self.position_embeddings is not None:
            position_embeddings = rearrange_channels(self.position_embeddings, True, channels_first)
            position_embeddings = repeat(position_embeddings, "1 ... -> b ...", b=b)
        else:
            if isinstance(grid_size, int):
                grid_size = (grid_size, grid_size, grid_size)

            if crop_offsets is None:
                cache_key = (dim, grid_size, None)
                if cache_key not in self.position_embeddings_cache:
                    self.position_embeddings_cache[cache_key] = get_absolute_position_embeddings_3d(
                        dim, grid_size, channels_first=channels_first
                    )
                position_embeddings = self.position_embeddings_cache[cache_key]
                position_embeddings = repeat(position_embeddings, "1 ... -> b ...", b=b)
            else:
                if crop_offsets.ndim == 1:
                    crop_offsets = crop_offsets.unsqueeze(0)

                position_embeddings = []
                for crop_offset in crop_offsets:
                    position_embeddings.append(
                        get_absolute_position_embeddings_3d(
                            dim, grid_size, crop_offset=crop_offset.tolist(), channels_first=channels_first
                        )
                    )
                position_embeddings = torch.cat(position_embeddings, dim=0)
            position_embeddings = position_embeddings.to(x.device)
        # (b, [dim], d, h, w, [dim])

        # Incorporate spacing information
        if spacings is not None:
            assert spacings.shape == (b, 3), "spacings must be of shape (batch_size, 3)"
            assert dim % 3 == 0, "dim must be divisible by 3"
            # (b, 3)
            spacings = repeat(spacings, "b three -> b (three dim_by_three) 1 1 1", three=3, dim_by_three=dim // 3)
            # (b, dim, 1, 1, 1)
            spacings = rearrange_channels(spacings, True, channels_first)
            # (b, [dim], 1, 1, 1, [dim])

            position_embeddings = position_embeddings * spacings.to(position_embeddings.device)
            # (b, [dim], d, h, w, [dim])

        if embedding_type == "add":
            x = x + position_embeddings
        elif embedding_type == "concat":
            x = torch.cat([x, position_embeddings], dim=1 if channels_first else -1)
        else:
            raise NotImplementedError("Only 'add' and 'concat' are supported for embedding_type")

        return x

# %% ../../nbs/layers/02_embeddings.ipynb 17
def get_specific_sinusoidal_embeddings_1d(dim: int, indices: torch.Tensor) -> torch.Tensor:
    """Get 1D sinusoidal position embeddings for specific indices.

    Args:
        dim: Embedding dimension. Must be divisible by 2.
        indices: Indices for which to get the embeddings. Shape: (length,).

    Returns:
        A tensor of shape (1, length, dim) containing the position embeddings.
    """
    if dim % 2 != 0:
        raise ValueError("dim must be divisible by 2")

    # Create frequency bands
    omega = torch.arange(dim // 2, dtype=torch.float32, device=indices.device)
    omega /= dim / 2.0
    omega = 1.0 / (10000**omega)
    # (dim // 2)

    # Outer product of positions / timesteps and frequencies
    out = torch.einsum("n,d->nd", indices, omega)
    # (length, dim//2)

    # Apply sin and cos functions
    emb_sin = torch.sin(out)
    emb_cos = torch.cos(out)

    # Interleave sin and cos embeddings
    embeddings = torch.stack([emb_sin, emb_cos], dim=2)
    embeddings = embeddings.flatten(1)
    # (length, dim)

    # Reshape to expected output format
    embeddings = rearrange(embeddings, "length dim -> 1 length dim").contiguous()
    # (1, length, dim)

    return embeddings


def get_sinusoidal_embeddings_1d(dim: int, length: int, device=torch.device("cpu")) -> torch.Tensor:
    """Get 1D sinusoidal position embeddings.

    Args:
        dim: Embedding dimension. Must be divisible by 2.
        length: Length of the sequence.
        device: Device to create the embeddings on.

    Returns:
        A tensor of shape (1, length, dim) containing the position embeddings.
    """
    # Create position / timestep indices
    indices = torch.arange(length, dtype=torch.int32, device=device)
    # (length,)

    return get_specific_sinusoidal_embeddings_1d(dim, indices)


get_timestep_embeddings_1d = get_specific_sinusoidal_embeddings_1d
get_all_timestep_embeddings_1d = get_sinusoidal_embeddings_1d
get_absolute_position_embeddings_1d = get_sinusoidal_embeddings_1d

# %% ../../nbs/layers/02_embeddings.ipynb 19
@populate_docstring
class AbsolutePositionEmbeddings1D(nn.Module):
    """1D Absolute Position Embeddings. May or may not learnable. These have to be applied on the input manually and
    cannot be passed to attention layers directly. {CLASS_DESCRIPTION_1D_DOC}"""

    @populate_docstring
    def __init__(self, config: AbsolutePositionEmbeddings1DConfig = {}, **kwargs):
        """Initialize AbsolutePositionEmbeddings1D.

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

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

        dim = self.config.dim
        length = self.config.length
        learnable = self.config.learnable

        self.position_embeddings_cache = {}
        self.position_embeddings = None
        if dim is not None and length is not None:
            self.position_embeddings = nn.Parameter(
                get_absolute_position_embeddings_1d(dim, length), requires_grad=learnable
            )

    @populate_docstring
    def forward(
        self,
        x: torch.Tensor,
        embedding_type: Literal["add", "concat"] = "add",
    ) -> torch.Tensor:
        """Apply absolute position embeddings to the input tensor.

        Args:
            x: {INPUT_1D_DOC}
            embedding_type: Type of embedding to apply. 'add' to add the position embeddings to the input,
                'concat' to concatenate them along the last dimension.

        Returns:
            {OUTPUT_1D_DOC}
        """
        assert x.ndim == 3, "Input tensor must be of shape (b, length, dim)"
        # Check if sufficient information has been provided
        if self.position_embeddings is None:
            dim = self.config.dim
            if dim is None:
                dim = x.shape[2]
            length = self.config.length
            if length is None:
                length = x.shape[1]
        else:
            dim = self.config.dim
            length = self.config.length

        # Estimate batch size
        b = x.shape[0]

        # Get position embeddings, adjust based on crop offsets if applicable
        if self.position_embeddings is not None:
            position_embeddings = self.position_embeddings
            position_embeddings = repeat(position_embeddings, "1 l d-> b l d", b=b)
        else:
            cache_key = (dim, length)
            if cache_key not in self.position_embeddings_cache:
                self.position_embeddings_cache[cache_key] = get_absolute_position_embeddings_1d(dim, length)
            position_embeddings = self.position_embeddings_cache[cache_key]
            position_embeddings = repeat(position_embeddings, "1 l d -> b l d", b=b).to(x.device)
        # (b, length, dim)

        if embedding_type == "add":
            x = x + position_embeddings
        elif embedding_type == "concat":
            x = torch.cat([x, position_embeddings], dim=1)
        else:
            raise NotImplementedError("Only 'add' and 'concat' are supported for embedding_type")

        return x

# %% ../../nbs/layers/02_embeddings.ipynb 22
@populate_docstring
class PatchEmbeddings3D(CNNBlock3D):
    """3D Patch Embeddings using a convolutional layer. {CLASS_DESCRIPTION_3D_DOC}"""

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

        Args:
            config: {CONFIG_INSTANCE_DOC}
            checkpointing_level: {CHECKPOINTING_LEVEL_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        self.config = PatchEmbeddings3DConfig.model_validate(config | kwargs)
        config = self.config.model_dump() | {
            "kernel_size": self.config.get("patch_size"),
            "stride": self.config.get("patch_size"),
            "padding": 0,
            "out_channels": self.config.get("dim"),
        }
        super().__init__(config, checkpointing_level, **kwargs)
