chuchichaestli.models.autoencoder

Autoencoder implementations and utilities.

Modules:

Name	Description
autoencoder	A highly-customizable autoencoder implementation.
dcae	A flexible deep-compression autoencoder implementation.
decoder	Decoder modules for autoencoders.
encoder	Encoder modules for autoencoders (AEs).
litevae	A lightweight wavelet-domain variational autoencoder.
traits	Traits the two components of an autoencoder implementation.
vae	A flexible variational autoencoder implementation.
vqvae	A flexible vector-quantized variational autoencoder implementation.

Classes:

Name	Description
Autoencoder	Flexible autoencoder implementation.
DCAE	Flexible deep-compression autoencoder implementation.
DCDecoder	Decoding component of a deep-compression autoencoder.
DCEncoder	Encoding component of a deep-compression autoencoder.
Decoder	Flexible decoder implementation for autoencoders.
DecoderLike	Interface an autoencoder expects from its decoding component.
Encoder	Flexible encoder implementation for autoencoders.
EncoderLike	Interface an autoencoder expects from its encoding component.
LiteVAE	Lightweight wavelet-domain variational autoencoder.
LiteVAEDecoder	Decoding component of a LiteVAE.
LiteVAEEncoder	Wavelet-domain encoding component of a LiteVAE.
LiteVAE_B	LiteVAE with the base published encoder, about 7M parameters (default).
LiteVAE_L	LiteVAE with the large published encoder, about 42M parameters.
LiteVAE_M	LiteVAE with the medium published encoder, about 32M parameters.
LiteVAE_S	LiteVAE with the small published encoder, about 1M parameters.
VAE	Flexible variational autoencoder implementation.
VAEDecoder	Decoding component of a variational autoencoder.
VAEEncoder	Encoding component of a variational autoencoder.
VQVAE	Flexible vector-quantized variational autoencoder implementation.
VectorQuantizer	Vector Quantizer for VQVAE.

Autoencoder

Autoencoder(
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = True,
    latent_deproj: nn.Module | bool = True,
)

Bases: nn.Module

Flexible autoencoder implementation.

The architecture consists of an encoder-decoder structure. The encoder chains several residual and downsampling blocks. Each downsampling block separates the encoder into spatially hierarchical levels. The encoder ends in bottleneck blocks (optionally including attention blocks and a convolutional layer) and projects the input into latent space. The decoder is built with residual convolutional and upsampling blocks, and expands from the latent space to the image domain.

Both components are constructed externally and passed in, which lets a configuration framework instantiate them as separate groups. Use the build constructor to assemble a model from architecture arguments instead.

Attributes:

Name	Type	Description
encoder_cls	type	Encoder class that build instantiates.
decoder_cls	type	Decoder class that build instantiates.

Assemble an autoencoder from an encoding and a decoding component.

Parameters:

Name	Type	Description	Default
encoder	EncoderLike	Encoding component, mapping the input to latent space.	required
decoder	DecoderLike	Decoding component, expanding the latent space to the output.	required
latent_proj	nn.Module | bool	Projection between encoder and latent space; True builds a pointwise convolution, False omits it, and a module is used as given.	True
latent_deproj	nn.Module | bool	Projection between latent space and decoder, with the same options as latent_proj.	True

Methods:

Name	Description
build	Build a model from architecture arguments, components included.
compute_latent_shape	Compute the shape of the latent space.
decode	Decode the input.
encode	Encode the input.
forward	Forward pass through the model, i.e. encode and decode.
projection	Pointwise convolution matching the component it serves.

Source code in src/chuchichaestli/models/autoencoder/autoencoder.py

def __init__(
    self,
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = True,
    latent_deproj: nn.Module | bool = True,
):
    """Assemble an autoencoder from an encoding and a decoding component.

    Args:
        encoder: Encoding component, mapping the input to latent space.
        decoder: Decoding component, expanding the latent space to the output.
        latent_proj: Projection between encoder and latent space; `True`
            builds a pointwise convolution, `False` omits it, and a module
            is used as given.
        latent_deproj: Projection between latent space and decoder, with the
            same options as `latent_proj`.
    """
    super().__init__()

    self.encoder = encoder
    self.decoder = decoder
    if latent_proj is True:
        latent_proj = self.projection(
            encoder, encoder.out_channels, encoder.out_channels
        )
    if latent_deproj is True:
        latent_deproj = self.projection(decoder, self.latent_dim, self.latent_dim)
    self.latent_proj = latent_proj if isinstance(latent_proj, nn.Module) else None
    self.latent_deproj = (
        latent_deproj if isinstance(latent_deproj, nn.Module) else None
    )
    self._check_components()

channel_mults property

channel_mults: int

Total channel multiplication across the encoder levels.

f_comp property

f_comp: int

Spatial compression factor of the encoder (number of spatial downsampling layers).

f_exp property

f_exp: int

Spatial expansion factor of the decoder (number of spatial upsampling layers).

latent_dim property

latent_dim: int

Latent channel dimension.

levels property

levels: tuple[int, int]

Number of stages in the encoder and decoder.

build classmethod

build(
    dimensions: int = 2,
    in_channels: int = 1,
    out_channels: int = 1,
    latent_dim: int = 4,
    decoder_n_channels: int | None = None,
    res_act_fn: ActivationTypes = "silu",
    res_dropout: float = 0.0,
    res_norm_type: NormTypes = "group",
    res_groups: int = 8,
    res_kernel_size: int = 3,
    attn_head_dim: int = 32,
    attn_n_heads: int = 1,
    attn_dropout_p: float = 0.0,
    attn_norm_type: NormTypes = "group",
    attn_groups: int = 32,
    attn_kernel_size: int = 1,
    attn_scales: Sequence[int] = (5,),
    context_args: dict = {},
    local_args: dict = {},
    encoder_args: dict = {},
    decoder_args: dict = {},
    **kwargs,
) -> Autoencoder

Build a model from architecture arguments, components included.

The res_* and attn_* arguments are shared by both components; anything specific to one component goes into encoder_args or decoder_args, whose keys are the parameter names of Encoder and Decoder. Keys left out fall back to the defaults of encoder_cls and decoder_cls.

Parameters:

Name	Type	Description	Default
dimensions	int	Number of dimensions for the model.	2
in_channels	int	Number of input channels.	1
out_channels	int	Number of output channels.	1
latent_dim	int	Number of channels in the latent space.	4
decoder_n_channels	int | None	Number of channels the decoder starts from; the encoder's bottleneck width if omitted, which is what mirrored components want. Give it when the two halves are sized independently.	None
res_act_fn	ActivationTypes	Activation function for the residual blocks (see chuchichaestli.models.activations for details).	'silu'
res_dropout	float	Dropout rate for the residual blocks.	0.0
res_norm_type	NormTypes	Normalization type for the residual block (see chuchichaestli.models.norm for details).	'group'
res_groups	int	Number of groups for the residual block normalization (if group norm).	8
res_kernel_size	int	Kernel size for the residual blocks.	3
attn_head_dim	int	Dimension of the attention heads.	32
attn_n_heads	int	Number of attention heads.	1
attn_dropout_p	float	Dropout probability of the scaled dot product attention.	0.0
attn_norm_type	NormTypes	Normalization type for the convolutional attention block (see chuchichaestli.models.norm for details).	'group'
attn_groups	int	Number of groups for the convolutional attention block normalization (if attn_norm_type is "group").	32
attn_kernel_size	int	Kernel size for the convolutional attention block.	1
attn_scales	Sequence[int]	Scales for the multi-scale attention block.	(5,)
context_args	dict	Keyword arguments for the context block in a transformer module.	{}
local_args	dict	Keyword arguments for the local block in a transformer module.	{}
encoder_args	dict	Arguments for the encoding component, overriding its defaults. Its res_args and attn_args override the shared values above.	{}
decoder_args	dict	Arguments for the decoding component, with the same options.	{}
kwargs		Keyword arguments for the constructor, such as latent_proj.	{}

Returns:

Type	Description
Autoencoder	An assembled model of the class this was called on.

Source code in src/chuchichaestli/models/autoencoder/autoencoder.py

@classmethod
def build(
    cls,
    dimensions: int = 2,
    in_channels: int = 1,
    out_channels: int = 1,
    latent_dim: int = 4,
    decoder_n_channels: int | None = None,
    res_act_fn: ActivationTypes = "silu",
    res_dropout: float = 0.0,
    res_norm_type: NormTypes = "group",
    res_groups: int = 8,
    res_kernel_size: int = 3,
    attn_head_dim: int = 32,
    attn_n_heads: int = 1,
    attn_dropout_p: float = 0.0,
    attn_norm_type: NormTypes = "group",
    attn_groups: int = 32,
    attn_kernel_size: int = 1,
    attn_scales: Sequence[int] = (5,),
    context_args: dict = {},
    local_args: dict = {},
    encoder_args: dict = {},
    decoder_args: dict = {},
    **kwargs,
) -> "Autoencoder":
    """Build a model from architecture arguments, components included.

    The `res_*` and `attn_*` arguments are shared by both components; anything
    specific to one component goes into `encoder_args` or `decoder_args`, whose
    keys are the parameter names of `Encoder` and `Decoder`. Keys left out
    fall back to the defaults of `encoder_cls` and `decoder_cls`.

    Args:
        dimensions: Number of dimensions for the model.
        in_channels: Number of input channels.
        out_channels: Number of output channels.
        latent_dim: Number of channels in the latent space.
        decoder_n_channels: Number of channels the decoder starts from; the
            encoder's bottleneck width if omitted, which is what mirrored
            components want. Give it when the two halves are sized
            independently.
        res_act_fn: Activation function for the residual blocks
            (see `chuchichaestli.models.activations` for details).
        res_dropout: Dropout rate for the residual blocks.
        res_norm_type: Normalization type for the residual block
            (see `chuchichaestli.models.norm` for details).
        res_groups: Number of groups for the residual block normalization (if group norm).
        res_kernel_size: Kernel size for the residual blocks.
        attn_head_dim: Dimension of the attention heads.
        attn_n_heads: Number of attention heads.
        attn_dropout_p: Dropout probability of the scaled dot product attention.
        attn_norm_type: Normalization type for the convolutional attention block
            (see `chuchichaestli.models.norm` for details).
        attn_groups: Number of groups for the convolutional attention block normalization
            (if `attn_norm_type` is `"group"`).
        attn_kernel_size: Kernel size for the convolutional attention block.
        attn_scales: Scales for the multi-scale attention block.
        context_args: Keyword arguments for the context block in a transformer module.
        local_args: Keyword arguments for the local block in a transformer module.
        encoder_args: Arguments for the encoding component, overriding its defaults.
            Its `res_args` and `attn_args` override the shared values above.
        decoder_args: Arguments for the decoding component, with the same options.
        kwargs: Keyword arguments for the constructor, such as `latent_proj`.

    Returns:
        An assembled model of the class this was called on.
    """
    _reject_sequences(
        res_act_fn=res_act_fn,
        res_dropout=res_dropout,
        res_groups=res_groups,
        res_norm_type=res_norm_type,
        res_kernel_size=res_kernel_size,
        attn_head_dim=attn_head_dim,
        attn_n_heads=attn_n_heads,
        attn_dropout_p=attn_dropout_p,
        attn_groups=attn_groups,
        attn_kernel_size=attn_kernel_size,
    )

    res_args = {
        "res_act_fn": res_act_fn,
        "res_dropout": res_dropout,
        "res_groups": res_groups,
        "res_norm_type": res_norm_type,
        "res_kernel_size": res_kernel_size,
    }

    attn_args = {
        "n_heads": attn_n_heads,
        "head_dim": attn_head_dim,
        "dropout_p": attn_dropout_p,
        "norm_type": attn_norm_type,
        "kernel_size": attn_kernel_size,
        "groups": attn_groups,
        "scales": attn_scales,
        "context_args": context_args,
        "local_args": local_args,
    }

    enc_args = _merge_shared_args(
        encoder_args, res_args, attn_args, ENCODER_RESERVED, "encoder_args"
    )
    dec_args = _merge_shared_args(
        decoder_args, res_args, attn_args, DECODER_RESERVED, "decoder_args"
    )
    encoder = cls.encoder_cls(
        dimensions=dimensions,
        in_channels=in_channels,
        out_channels=latent_dim,
        **enc_args,
    )
    decoder = cls.decoder_cls(
        dimensions=dimensions,
        in_channels=encoder.latent_channels,
        n_channels=(
            encoder.bottleneck_channels
            if decoder_n_channels is None
            else decoder_n_channels
        ),
        out_channels=out_channels,
        **dec_args,
    )
    return cls(encoder=encoder, decoder=decoder, **kwargs)

compute_latent_shape

compute_latent_shape(
    input_shape: tuple[int, ...], no_batch_dim: bool = False
)

Compute the shape of the latent space.

Source code in src/chuchichaestli/models/autoencoder/autoencoder.py

def compute_latent_shape(
    self, input_shape: tuple[int, ...], no_batch_dim: bool = False
):
    """Compute the shape of the latent space."""
    batch_dim = input_shape[0] if not no_batch_dim else None
    spatial_dims = tuple(dim // self.f_comp for dim in input_shape[2:])
    if batch_dim is None:
        shape = (self.latent_dim, *spatial_dims)
    else:
        shape = (batch_dim, self.latent_dim, *spatial_dims)
    return shape

decode

decode(z: torch.Tensor) -> torch.Tensor

Decode the input.

Parameters:

Name	Type	Description	Default
z	torch.Tensor	Input latent tensor.	required

Returns:

Type	Description
torch.Tensor	Image reconstructed from latent code.

Source code in src/chuchichaestli/models/autoencoder/autoencoder.py

def decode(self, z: torch.Tensor) -> torch.Tensor:
    """Decode the input.

    Args:
        z: Input latent tensor.

    Returns:
        Image reconstructed from latent code.
    """
    z = self.latent_deproj(z) if self.latent_deproj is not None else z
    return self.decoder(z)

encode

encode(x: torch.Tensor) -> torch.Tensor

Encode the input.

Parameters:

Name	Type	Description	Default
x	torch.Tensor	Input tensor.	required
eps		Small constant value for numerical stability.	required

Returns:

Type	Description
torch.Tensor	Multivariate normal posterior distribution

Source code in src/chuchichaestli/models/autoencoder/autoencoder.py

def encode(self, x: torch.Tensor) -> torch.Tensor:
    """Encode the input.

    Args:
        x: Input tensor.
        eps: Small constant value for numerical stability.

    Returns:
        Multivariate normal posterior distribution
    """
    z = self.encoder(x)
    z = self.latent_proj(z) if self.latent_proj is not None else z
    return z

forward

forward(x: torch.Tensor) -> torch.Tensor

Forward pass through the model, i.e. encode and decode.

Source code in src/chuchichaestli/models/autoencoder/autoencoder.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass through the model, i.e. encode and decode."""
    code = self.encode(x)
    return self.decode(code)

projection staticmethod

projection(
    component: EncoderLike | DecoderLike,
    in_channels: int,
    out_channels: int,
) -> nn.Module

Pointwise convolution matching the component it serves.

Device and dtype are read from the component's first parameter or buffer, so a component placed before the model is assembled keeps the projection alongside it; a component holding neither leaves both to the torch defaults.

Parameters:

Name	Type	Description	Default
component	EncoderLike | DecoderLike	Encoding or decoding component the projection is attached to.	required
in_channels	int	Number of input channels.	required
out_channels	int	Number of output channels.	required

Source code in src/chuchichaestli/models/autoencoder/autoencoder.py

@staticmethod
def projection(
    component: EncoderLike | DecoderLike, in_channels: int, out_channels: int
) -> nn.Module:
    """Pointwise convolution matching the component it serves.

    Device and dtype are read from the component's first parameter or
    buffer, so a component placed before the model is assembled keeps the
    projection alongside it; a component holding neither leaves both to the
    torch defaults.

    Args:
        component: Encoding or decoding component the projection is attached to.
        in_channels: Number of input channels.
        out_channels: Number of output channels.
    """
    ref = next(chain(component.parameters(), component.buffers()), None)
    return DIM_TO_CONV_MAP[component.dimensions](
        in_channels,
        out_channels,
        kernel_size=1,
        stride=1,
        padding="same",
        device=None if ref is None else ref.device,
        dtype=None if ref is None else ref.dtype,
    )

DCAE

DCAE(
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
)

Bases: Autoencoder

Flexible deep-compression autoencoder implementation.

The architecture consists of an encoder-decoder structure. The encoder chains several residual and downsampling blocks. Each downsampling block separates the encoder into spatially hierarchical levels. The encoder ends in bottleneck blocks (optionally including attention blocks and a convolutional layer) and projects the input into latent space. The decoder is built with residual convolutional, efficient ViT blocks, and upsampling blocks, and expands from the latent space to the image domain.

The compression happens in the blocks themselves, so no latent projection is used by default.

Attributes:

Name	Type	Description
encoder_cls	type	Encoder class that build instantiates.
decoder_cls	type	Decoder class that build instantiates.

Assemble a deep-compression autoencoder from its two components.

Parameters:

Name	Type	Description	Default
encoder	EncoderLike	Encoding component, mapping the input to latent space.	required
decoder	DecoderLike	Decoding component, expanding the latent space to the output.	required
latent_proj	nn.Module | bool	Projection between encoder and latent space (see Autoencoder).	False
latent_deproj	nn.Module | bool	Projection between latent space and decoder.	False

Methods:

Name	Description
build	Build a deep-compression autoencoder from architecture arguments.

Source code in src/chuchichaestli/models/autoencoder/dcae.py

def __init__(
    self,
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
):
    """Assemble a deep-compression autoencoder from its two components.

    Args:
        encoder: Encoding component, mapping the input to latent space.
        decoder: Decoding component, expanding the latent space to the output.
        latent_proj: Projection between encoder and latent space (see `Autoencoder`).
        latent_deproj: Projection between latent space and decoder.
    """
    super().__init__(encoder, decoder, latent_proj, latent_deproj)

build classmethod

build(
    latent_dim: int = 32,
    res_norm_type: NormTypes = "rms",
    attn_norm_type: NormTypes
    | Sequence[NormTypes]
    | None = "rms",
    context_args: dict = DC_ATTN_ARGS["context_args"],
    local_args: dict = DC_ATTN_ARGS["local_args"],
    **kwargs,
) -> DCAE

Build a deep-compression autoencoder from architecture arguments.

Only the arguments whose defaults differ from Autoencoder.build are named here; the components supply every structural default.

Parameters:

Name	Type	Description	Default
latent_dim	int	Number of channels in the latent space.	32
res_norm_type	NormTypes	Normalization type for the residual blocks.	'rms'
attn_norm_type	NormTypes | Sequence[NormTypes] | None	Normalization type for the attention blocks.	'rms'
context_args	dict	Keyword arguments for the context block in a transformer module.	DC_ATTN_ARGS['context_args']
local_args	dict	Keyword arguments for the local block in a transformer module.	DC_ATTN_ARGS['local_args']
kwargs		Further architecture arguments, see Autoencoder.build.	{}

Returns:

Type	Description
DCAE	An assembled model of the class this was called on.

Source code in src/chuchichaestli/models/autoencoder/dcae.py

@classmethod
def build(
    cls,
    latent_dim: int = 32,
    res_norm_type: NormTypes = "rms",
    attn_norm_type: NormTypes | Sequence[NormTypes] | None = "rms",
    context_args: dict = DC_ATTN_ARGS["context_args"],
    local_args: dict = DC_ATTN_ARGS["local_args"],
    **kwargs,
) -> "DCAE":
    """Build a deep-compression autoencoder from architecture arguments.

    Only the arguments whose defaults differ from `Autoencoder.build` are
    named here; the components supply every structural default.

    Args:
        latent_dim: Number of channels in the latent space.
        res_norm_type: Normalization type for the residual blocks.
        attn_norm_type: Normalization type for the attention blocks.
        context_args: Keyword arguments for the context block in a transformer module.
        local_args: Keyword arguments for the local block in a transformer module.
        kwargs: Further architecture arguments, see `Autoencoder.build`.

    Returns:
        An assembled model of the class this was called on.
    """
    return super().build(
        latent_dim=latent_dim,
        res_norm_type=res_norm_type,
        attn_norm_type=attn_norm_type,
        context_args=context_args,
        local_args=local_args,
        **kwargs,
    )

DCDecoder

DCDecoder(
    dimensions: int = 2,
    in_channels: int = 32,
    n_channels: int = 1024,
    out_channels: int = 1,
    in_block_type: DecoderInBlockTypes = "DCDecoderInBlock",
    mid_block_types: Sequence[
        AutoencoderMidBlockTypes
    ] = (),
    up_block_types: Sequence[AutoencoderUpBlockTypes] = (
        "EfficientViTBlock",
    )
    * 3
    + ("DCAutoencoderUpBlock",) * 3,
    block_out_channel_mults: Sequence[int] = (
        1,
        2,
        1,
        2,
        2,
    ),
    num_layers_per_block: int | Sequence[int] = 3,
    upsample_type: UpsampleTypes = "UpsampleShuffle",
    act_fn: ActivationTypes = "relu",
    norm_type: NormTypes = "rms",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    attn_args: dict = {},
    in_shortcut: bool = True,
)

Bases: Decoder

Decoding component of a deep-compression autoencoder.

Expands with channel-shuffling upsamplers, mirroring DCEncoder; see Decoder for what each argument does. Blocks are listed in order of data flow, deepest level first, so attention precedes the convolutional blocks.

Constructor.

Parameters:

Name	Type	Description	Default
dimensions	int	Number of dimensions.	2
in_channels	int	Number of input channels (latent space).	32
n_channels	int	Number of channels for the first block.	1024
out_channels	int	Number of output channels.	1
in_block_type	DecoderInBlockTypes	Type of block for input (latent space).	'DCDecoderInBlock'
mid_block_types	Sequence[AutoencoderMidBlockTypes]	Type of blocks to use after the input.	()
up_block_types	Sequence[AutoencoderUpBlockTypes]	Type of up blocks to use for each level.	('EfficientViTBlock',) * 3 + ('DCAutoencoderUpBlock',) * 3
block_out_channel_mults	Sequence[int]	Divisor for output channels of each block.	(1, 2, 1, 2, 2)
num_layers_per_block	int | Sequence[int]	Number of blocks per level.	3
upsample_type	UpsampleTypes	Type of upsampling block, per level.	'UpsampleShuffle'
act_fn	ActivationTypes	Activation function for the output layers.	'relu'
norm_type	NormTypes	Normalization type for the output layer.	'rms'
num_groups	int	Number of groups for normalization in the output layer.	8
kernel_size	int	Kernel size for the output convolution.	3
res_args	dict	Arguments for residual blocks, overriding DC_RES_ARGS.	{}
attn_args	dict	Arguments for attention blocks, overriding DC_ATTN_ARGS.	{}
in_shortcut	bool	Whether to use a shortcut for the input block.	True

Source code in src/chuchichaestli/models/autoencoder/dcae.py

def __init__(
    self,
    dimensions: int = 2,
    in_channels: int = 32,
    n_channels: int = 1024,
    out_channels: int = 1,
    in_block_type: DecoderInBlockTypes = "DCDecoderInBlock",
    mid_block_types: Sequence[AutoencoderMidBlockTypes] = (),
    up_block_types: Sequence[AutoencoderUpBlockTypes] = ("EfficientViTBlock",) * 3
    + ("DCAutoencoderUpBlock",) * 3,
    block_out_channel_mults: Sequence[int] = (1, 2, 1, 2, 2),
    num_layers_per_block: int | Sequence[int] = 3,
    upsample_type: UpsampleTypes = "UpsampleShuffle",
    act_fn: ActivationTypes = "relu",
    norm_type: NormTypes = "rms",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    attn_args: dict = {},
    in_shortcut: bool = True,
):
    """Constructor.

    Args:
        dimensions: Number of dimensions.
        in_channels: Number of input channels (latent space).
        n_channels: Number of channels for the first block.
        out_channels: Number of output channels.
        in_block_type: Type of block for input (latent space).
        mid_block_types: Type of blocks to use after the input.
        up_block_types: Type of up blocks to use for each level.
        block_out_channel_mults: Divisor for output channels of each block.
        num_layers_per_block: Number of blocks per level.
        upsample_type: Type of upsampling block, per level.
        act_fn: Activation function for the output layers.
        norm_type: Normalization type for the output layer.
        num_groups: Number of groups for normalization in the output layer.
        kernel_size: Kernel size for the output convolution.
        res_args: Arguments for residual blocks, overriding `DC_RES_ARGS`.
        attn_args: Arguments for attention blocks, overriding `DC_ATTN_ARGS`.
        in_shortcut: Whether to use a shortcut for the input block.
    """
    super().__init__(
        dimensions=dimensions,
        in_channels=in_channels,
        n_channels=n_channels,
        out_channels=out_channels,
        in_block_type=in_block_type,
        mid_block_types=mid_block_types,
        up_block_types=up_block_types,
        block_out_channel_mults=block_out_channel_mults,
        num_layers_per_block=num_layers_per_block,
        upsample_type=upsample_type,
        act_fn=act_fn,
        norm_type=norm_type,
        num_groups=num_groups,
        kernel_size=kernel_size,
        res_args={**DC_RES_ARGS, **res_args},
        attn_args={**DC_ATTN_ARGS, **attn_args},
        in_shortcut=in_shortcut,
    )

DCEncoder

DCEncoder(
    dimensions: int = 2,
    in_channels: int = 1,
    n_channels: int = 128,
    out_channels: int = 32,
    down_block_types: Sequence[
        AutoencoderDownBlockTypes
    ] = ("DCAutoencoderDownBlock",) * 3
    + ("EfficientViTBlock",) * 3,
    block_out_channel_mults: Sequence[int] = (
        2,
        2,
        1,
        2,
        1,
    ),
    num_layers_per_block: int | Sequence[int] = (
        2,
        2,
        2,
        3,
        3,
        3,
    ),
    mid_block_types: Sequence[
        AutoencoderMidBlockTypes
    ] = (),
    out_block_type: EncoderOutBlockTypes = "DCEncoderOutBlock",
    downsample_type: DownsampleTypes = "DownsampleUnshuffle",
    act_fn: ActivationTypes = "silu",
    norm_type: NormTypes = "rms",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    attn_args: dict = {},
    out_shortcut: bool = True,
)

Bases: Encoder

Encoding component of a deep-compression autoencoder.

Compresses with channel-shuffling downsamplers and efficient attention at the deeper levels; see Encoder for what each argument does.

Constructor.

Parameters:

Name	Type	Description	Default
dimensions	int	Number of dimensions.	2
in_channels	int	Number of input channels.	1
n_channels	int	Number of channels in the hidden layer.	128
out_channels	int	Number of output channels (latent space).	32
down_block_types	Sequence[AutoencoderDownBlockTypes]	Type of down blocks to use for each level.	('DCAutoencoderDownBlock',) * 3 + ('EfficientViTBlock',) * 3
block_out_channel_mults	Sequence[int]	Multiplier for output channels of each block.	(2, 2, 1, 2, 1)
num_layers_per_block	int | Sequence[int]	Number of blocks per level.	(2, 2, 2, 3, 3, 3)
mid_block_types	Sequence[AutoencoderMidBlockTypes]	Type of blocks to use before the output.	()
out_block_type	EncoderOutBlockTypes	Type of block for output (latent space).	'DCEncoderOutBlock'
downsample_type	DownsampleTypes	Type of downsampling block, per level.	'DownsampleUnshuffle'
act_fn	ActivationTypes	Activation function for the output layers.	'silu'
norm_type	NormTypes	Normalization type for the output layer.	'rms'
num_groups	int	Number of groups for normalization in the output layer.	8
kernel_size	int	Kernel size for the output convolution.	3
res_args	dict	Arguments for residual blocks, overriding DC_RES_ARGS.	{}
attn_args	dict	Arguments for attention blocks, overriding DC_ATTN_ARGS.	{}
out_shortcut	bool	Whether to use a shortcut for the output block.	True

Source code in src/chuchichaestli/models/autoencoder/dcae.py

def __init__(
    self,
    dimensions: int = 2,
    in_channels: int = 1,
    n_channels: int = 128,
    out_channels: int = 32,
    down_block_types: Sequence[AutoencoderDownBlockTypes] = (
        "DCAutoencoderDownBlock",
    )
    * 3
    + ("EfficientViTBlock",) * 3,
    block_out_channel_mults: Sequence[int] = (2, 2, 1, 2, 1),
    num_layers_per_block: int | Sequence[int] = (2, 2, 2, 3, 3, 3),
    mid_block_types: Sequence[AutoencoderMidBlockTypes] = (),
    out_block_type: EncoderOutBlockTypes = "DCEncoderOutBlock",
    downsample_type: DownsampleTypes = "DownsampleUnshuffle",
    act_fn: ActivationTypes = "silu",
    norm_type: NormTypes = "rms",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    attn_args: dict = {},
    out_shortcut: bool = True,
):
    """Constructor.

    Args:
        dimensions: Number of dimensions.
        in_channels: Number of input channels.
        n_channels: Number of channels in the hidden layer.
        out_channels: Number of output channels (latent space).
        down_block_types: Type of down blocks to use for each level.
        block_out_channel_mults: Multiplier for output channels of each block.
        num_layers_per_block: Number of blocks per level.
        mid_block_types: Type of blocks to use before the output.
        out_block_type: Type of block for output (latent space).
        downsample_type: Type of downsampling block, per level.
        act_fn: Activation function for the output layers.
        norm_type: Normalization type for the output layer.
        num_groups: Number of groups for normalization in the output layer.
        kernel_size: Kernel size for the output convolution.
        res_args: Arguments for residual blocks, overriding `DC_RES_ARGS`.
        attn_args: Arguments for attention blocks, overriding `DC_ATTN_ARGS`.
        out_shortcut: Whether to use a shortcut for the output block.
    """
    super().__init__(
        dimensions=dimensions,
        in_channels=in_channels,
        n_channels=n_channels,
        out_channels=out_channels,
        down_block_types=down_block_types,
        block_out_channel_mults=block_out_channel_mults,
        num_layers_per_block=num_layers_per_block,
        mid_block_types=mid_block_types,
        out_block_type=out_block_type,
        downsample_type=downsample_type,
        act_fn=act_fn,
        norm_type=norm_type,
        num_groups=num_groups,
        kernel_size=kernel_size,
        res_args={**DC_RES_ARGS, **res_args},
        attn_args={**DC_ATTN_ARGS, **attn_args},
        out_shortcut=out_shortcut,
    )

Decoder

Decoder(
    dimensions: int = 2,
    in_channels: int = 1,
    n_channels: int = 512,
    out_channels: int = 1,
    in_block_type: DecoderInBlockTypes = "DecoderInBlock",
    mid_block_types: Sequence[AutoencoderMidBlockTypes] = (
        "AutoencoderMidBlock",
        "AttnAutoencoderMidBlock",
    ),
    up_block_types: Sequence[AutoencoderUpBlockTypes] = (
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
    ),
    block_out_channel_mults: Sequence[int] = (1, 2, 2, 2),
    num_layers_per_block: int | Sequence[int] = 3,
    upsample_type: UpsampleTypes
    | Sequence[UpsampleTypes] = "UpsampleInterpolate",
    act_fn: ActivationTypes = "silu",
    norm_type: NormTypes = "group",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    attn_args: dict = {},
    in_shortcut: bool = False,
)

Bases: nn.Module

Flexible decoder implementation for autoencoders.

Decoder implementation.

Parameters:

Name	Type	Description	Default
dimensions	int	Number of dimensions.	2
in_channels	int	Number of input channels (latent space).	1
n_channels	int	Number of channels for first block.	512
out_channels	int	Number of output channels.	1
in_block_type	DecoderInBlockTypes	Type of block for output (latent space).	'DecoderInBlock'
mid_block_types	Sequence[AutoencoderMidBlockTypes]	Type of blocks to use before the output.	('AutoencoderMidBlock', 'AttnAutoencoderMidBlock')
up_block_types	Sequence[AutoencoderUpBlockTypes]	Type of up blocks to use for each level.	('AutoencoderUpBlock', 'AutoencoderUpBlock', 'AutoencoderUpBlock', 'AutoencoderUpBlock')
block_out_channel_mults	Sequence[int]	Multiplier for output channels of each block.	(1, 2, 2, 2)
num_layers_per_block	int | Sequence[int]	Number of blocks per level (blocks are repeated if >1).	3
upsample_type	UpsampleTypes | Sequence[UpsampleTypes]	Type of upsampling block, per level (see chuchichaestli.models.upsampling for details).	'UpsampleInterpolate'
act_fn	ActivationTypes	Activation function for the output layers (see chuchichaestli.models.activations for details).	'silu'
norm_type	NormTypes	Normalization type for the output layer.	'group'
num_groups	int	Number of groups for normalization in the output layer.	8
kernel_size	int	Kernel size for the output convolution.	3
res_args	dict	Arguments for residual blocks. Each value is either a single value or one per block position, in order of data flow.	{}
attn_args	dict	Arguments for attention blocks, per block position or per block that has attention. norm_type, scales, context_args and local_args are passed through unchanged.	{}
in_shortcut	bool	Whether to use a shortcut for the input block.	False

Methods:

Name	Description
forward	Decoding forward pass.

Attributes:

Name	Type	Description
f	int	Expansion factor of the decoder.

Source code in src/chuchichaestli/models/autoencoder/decoder.py

def __init__(
    self,
    dimensions: int = 2,
    in_channels: int = 1,
    n_channels: int = 512,
    out_channels: int = 1,
    in_block_type: DecoderInBlockTypes = "DecoderInBlock",
    mid_block_types: Sequence[AutoencoderMidBlockTypes] = (
        "AutoencoderMidBlock",
        "AttnAutoencoderMidBlock",
    ),
    up_block_types: Sequence[AutoencoderUpBlockTypes] = (
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
    ),
    block_out_channel_mults: Sequence[int] = (1, 2, 2, 2),
    num_layers_per_block: int | Sequence[int] = 3,
    upsample_type: UpsampleTypes | Sequence[UpsampleTypes] = "UpsampleInterpolate",
    act_fn: ActivationTypes = "silu",
    norm_type: NormTypes = "group",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    attn_args: dict = {},
    in_shortcut: bool = False,
):
    """Decoder implementation.

    Args:
        dimensions: Number of dimensions.
        in_channels: Number of input channels (latent space).
        n_channels: Number of channels for first block.
        out_channels: Number of output channels.
        in_block_type: Type of block for output (latent space).
        mid_block_types: Type of blocks to use before the output.
        up_block_types: Type of up blocks to use for each level.
        block_out_channel_mults: Multiplier for output channels of each block.
        num_layers_per_block: Number of blocks per level (blocks are repeated if `>1`).
        upsample_type: Type of upsampling block, per level (see
            `chuchichaestli.models.upsampling` for details).
        act_fn: Activation function for the output layers
            (see `chuchichaestli.models.activations` for details).
        norm_type: Normalization type for the output layer.
        num_groups: Number of groups for normalization in the output layer.
        kernel_size: Kernel size for the output convolution.
        res_args: Arguments for residual blocks. Each value is either a
            single value or one per block position, in order of data flow.
        attn_args: Arguments for attention blocks, per block position or per
            block that has attention. `norm_type`, `scales`, `context_args`
            and `local_args` are passed through unchanged.
        in_shortcut: Whether to use a shortcut for the input block.
    """
    super().__init__()

    block_out_channel_mults = tuple(block_out_channel_mults)
    if len(block_out_channel_mults) < len(up_block_types):
        block_out_channel_mults += (1,) * (
            len(up_block_types) - len(block_out_channel_mults)
        )
    elif len(block_out_channel_mults) > len(up_block_types):
        block_out_channel_mults = block_out_channel_mults[: len(up_block_types)]
    n_mults = len(block_out_channel_mults)
    self.channel_mults = prod(block_out_channel_mults)
    self.dimensions = dimensions
    self.in_channels = in_channels
    self.n_channels = n_channels
    self.out_channels = out_channels
    if isinstance(num_layers_per_block, int):
        num_layers_per_block = (num_layers_per_block,) * n_mults
    else:
        num_layers_per_block = tuple(num_layers_per_block)
        if len(num_layers_per_block) < len(up_block_types):
            num_layers_per_block += (num_layers_per_block[-1],) * (
                len(up_block_types) - len(num_layers_per_block)
            )

    # One sampler entry per level; the last governs channel bookkeeping only
    upsample_types = broadcast(
        upsample_type, n_mults, "upsample_type", None, f"[{n_mults} level(s)]"
    )
    upsample_clss = [
        require_cls(name, UPSAMPLE_FUNCTIONS, "upsampling type")
        for name in upsample_types
    ]

    # Block positions in order of data flow
    n_pos = n_mults + len(mid_block_types)
    path = f"[{len(mid_block_types)} mid block(s), {n_mults} level(s)]"
    has_attention = [
        block_type in ATTN_BLOCK_MAP
        for block_type in (*mid_block_types, *up_block_types)
    ]
    res_args = broadcast_kwargs(res_args, n_pos, context=path)
    attn_args = broadcast_kwargs(
        attn_args,
        n_pos,
        mask=has_attention,
        opaque=INTRA_BLOCK_ATTN_ARGS,
        context=path,
    )

    self.in_block = BLOCK_MAP[in_block_type](
        dimensions=dimensions,
        in_channels=in_channels,
        out_channels=n_channels,
        kernel_size=kernel_size,
    )

    self.mid_blocks = nn.ModuleList([])
    for j, mid_block_type in enumerate(mid_block_types):
        mid_block = BLOCK_MAP[mid_block_type](
            dimensions=dimensions,
            channels=n_channels,
            res_args=res_args[j],
            attn_args=attn_args[j],
        )
        self.mid_blocks.append(mid_block)

    self.up_blocks = nn.ModuleList([])
    ins = n_channels
    for i in range(n_mults):
        outs = ins
        if not upsample_clss[i].changes_channels:
            outs = int(ins // block_out_channel_mults[i])
        stage = nn.Sequential()
        for _ in range(num_layers_per_block[i]):
            up_block = BLOCK_MAP[up_block_types[i]](
                dimensions=dimensions,
                in_channels=ins,
                out_channels=outs,
                res_args=res_args[len(mid_block_types) + i],
                attn_args=attn_args[len(mid_block_types) + i],
            )
            stage.append(up_block)
            ins = outs
        self.up_blocks.append(stage)

        if i < n_mults - 1:
            if not upsample_clss[i].changes_channels:
                self.up_blocks.append(upsample_clss[i](dimensions, outs))
            else:
                self.up_blocks.append(
                    upsample_clss[i](
                        dimensions,
                        ins,
                        outs := int(ins // block_out_channel_mults[i]),
                    )
                )
                ins = outs
    self.levels = (len(self.up_blocks) + 1) // 2
    self.out_block = NormActConvBlock(
        dimensions=dimensions,
        in_channels=outs,
        out_channels=out_channels,
        act_fn=act_fn,
        norm_type=norm_type,
        num_groups=num_groups,
        kernel_size=kernel_size,
        stride=1,
        padding="same",
    )
    self.in_shortcut = in_shortcut
    self.shortcut_repeats = n_channels // in_channels

f property

f: int

Expansion factor of the decoder.

Raises:

Type	Description
ValueError	If a sampling block does not scale by a constant factor.

forward

forward(z: torch.Tensor) -> torch.Tensor

Decoding forward pass.

Source code in src/chuchichaestli/models/autoencoder/decoder.py

def forward(self, z: torch.Tensor) -> torch.Tensor:
    """Decoding forward pass."""
    if self.in_shortcut:
        shortcut = z.repeat_interleave(
            self.shortcut_repeats,
            dim=1,
            output_size=z.shape[1] * self.shortcut_repeats,
        )
        z = self.in_block(z) + shortcut
    else:
        z = self.in_block(z)
    for block in self.mid_blocks:
        z = block(z)
    for block in self.up_blocks:
        z = block(z)
    z = self.out_block(z)
    return z

DecoderLike

Bases: Protocol

Interface an autoencoder expects from its decoding component.

Any module exposing these members can be passed to Autoencoder; inheriting from Decoder is not required. Attributes absent from a custom component simply skip the consistency checks that read them.

Attributes:

Name	Type	Description
dimensions	int	Number of spatial dimensions.
in_channels	int	Number of latent channels consumed.
n_channels	int	Number of channels after the input block.
out_channels	int	Number of output channels.
channel_mults	int	Total channel division across all levels.
levels	int	Number of spatially hierarchical levels.

Methods:

Name	Description
__call__	Decode a latent tensor.
buffers	Iterate over the component's buffers.
parameters	Iterate over the component's parameters.

f property

f: int

Spatial expansion factor.

__call__

__call__(z: torch.Tensor) -> torch.Tensor

Decode a latent tensor.

Parameters:

Name	Type	Description	Default
z	torch.Tensor	Input latent tensor.	required

Source code in src/chuchichaestli/models/autoencoder/traits.py

def __call__(self, z: torch.Tensor) -> torch.Tensor:
    """Decode a latent tensor.

    Args:
        z: Input latent tensor.
    """
    ...

buffers

buffers() -> Iterator[torch.Tensor]

Iterate over the component's buffers.

Source code in src/chuchichaestli/models/autoencoder/traits.py

def buffers(self) -> Iterator[torch.Tensor]:
    """Iterate over the component's buffers."""
    ...

parameters

parameters() -> Iterator[nn.Parameter]

Iterate over the component's parameters.

Source code in src/chuchichaestli/models/autoencoder/traits.py

def parameters(self) -> Iterator[nn.Parameter]:
    """Iterate over the component's parameters."""
    ...

Encoder

Encoder(
    dimensions: int = 2,
    in_channels: int = 1,
    n_channels: int = 64,
    out_channels: int = 1,
    down_block_types: Sequence[
        AutoencoderDownBlockTypes
    ] = (
        "AutoencoderDownBlock",
        "AutoencoderDownBlock",
        "AutoencoderDownBlock",
        "AutoencoderDownBlock",
    ),
    block_out_channel_mults: Sequence[int] = (1, 2, 2, 2),
    num_layers_per_block: int | Sequence[int] = 2,
    mid_block_types: Sequence[AutoencoderMidBlockTypes] = (
        "AutoencoderMidBlock",
        "AttnAutoencoderMidBlock",
    ),
    out_block_type: EncoderOutBlockTypes = "EncoderOutBlock",
    downsample_type: DownsampleTypes
    | Sequence[DownsampleTypes] = "Downsample",
    act_fn: ActivationTypes = "silu",
    norm_type: NormTypes = "group",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    attn_args: dict = {},
    out_shortcut: bool = False,
)

Bases: nn.Module

Flexible encoder implementation for autoencoders.

Constructor.

Parameters:

Name	Type	Description	Default
dimensions	int	Number of dimensions.	2
in_channels	int	Number of input channels.	1
n_channels	int	Number of channels in the hidden layer.	64
out_channels	int	Number of output channels (latent space).	1
down_block_types	Sequence[AutoencoderDownBlockTypes]	Type of down blocks to use for each level.	('AutoencoderDownBlock', 'AutoencoderDownBlock', 'AutoencoderDownBlock', 'AutoencoderDownBlock')
block_out_channel_mults	Sequence[int]	Multiplier for output channels of each block.	(1, 2, 2, 2)
num_layers_per_block	int | Sequence[int]	Number of blocks per level (blocks are repeated if >1).	2
mid_block_types	Sequence[AutoencoderMidBlockTypes]	Type of blocks to use before the output.	('AutoencoderMidBlock', 'AttnAutoencoderMidBlock')
out_block_type	EncoderOutBlockTypes	Type of block for output (latent space).	'EncoderOutBlock'
downsample_type	DownsampleTypes | Sequence[DownsampleTypes]	Type of downsampling block, per level (see chuchichaestli.models.downsampling for details).	'Downsample'
act_fn	ActivationTypes	Activation function for the output layers (see chuchichaestli.models.activations for details).	'silu'
norm_type	NormTypes	Normalization type for the output layer.	'group'
num_groups	int	Number of groups for normalization in the output layer.	8
kernel_size	int	Kernel size for the output convolution.	3
res_args	dict	Arguments for residual blocks. Each value is either a single value or one per block position, in order of data flow.	{}
attn_args	dict	Arguments for attention blocks, per block position or per block that has attention. norm_type, scales, context_args and local_args are passed through unchanged.	{}
out_shortcut	bool	Whether to use a shortcut for the output block.	False

Methods:

Name	Description
forward	Forward pass.

Attributes:

Name	Type	Description
f	int	Compression factor of the encoder.

Source code in src/chuchichaestli/models/autoencoder/encoder.py

def __init__(
    self,
    dimensions: int = 2,
    in_channels: int = 1,
    n_channels: int = 64,
    out_channels: int = 1,
    down_block_types: Sequence[AutoencoderDownBlockTypes] = (
        "AutoencoderDownBlock",
        "AutoencoderDownBlock",
        "AutoencoderDownBlock",
        "AutoencoderDownBlock",
    ),
    block_out_channel_mults: Sequence[int] = (1, 2, 2, 2),
    num_layers_per_block: int | Sequence[int] = 2,
    mid_block_types: Sequence[AutoencoderMidBlockTypes] = (
        "AutoencoderMidBlock",
        "AttnAutoencoderMidBlock",
    ),
    out_block_type: EncoderOutBlockTypes = "EncoderOutBlock",
    downsample_type: DownsampleTypes | Sequence[DownsampleTypes] = "Downsample",
    act_fn: ActivationTypes = "silu",
    norm_type: NormTypes = "group",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    attn_args: dict = {},
    out_shortcut: bool = False,
):
    """Constructor.

    Args:
        dimensions: Number of dimensions.
        in_channels: Number of input channels.
        n_channels: Number of channels in the hidden layer.
        out_channels: Number of output channels (latent space).
        down_block_types: Type of down blocks to use for each level.
        block_out_channel_mults: Multiplier for output channels of each block.
        num_layers_per_block: Number of blocks per level (blocks are repeated if `>1`).
        mid_block_types: Type of blocks to use before the output.
        out_block_type: Type of block for output (latent space).
        downsample_type: Type of downsampling block, per level (see
            `chuchichaestli.models.downsampling` for details).
        act_fn: Activation function for the output layers
            (see `chuchichaestli.models.activations` for details).
        norm_type: Normalization type for the output layer.
        num_groups: Number of groups for normalization in the output layer.
        kernel_size: Kernel size for the output convolution.
        res_args: Arguments for residual blocks. Each value is either a
            single value or one per block position, in order of data flow.
        attn_args: Arguments for attention blocks, per block position or per
            block that has attention. `norm_type`, `scales`, `context_args`
            and `local_args` are passed through unchanged.
        out_shortcut: Whether to use a shortcut for the output block.
    """
    super().__init__()

    block_out_channel_mults = tuple(block_out_channel_mults)
    if len(block_out_channel_mults) < len(down_block_types):
        block_out_channel_mults += (1,) * (
            len(down_block_types) - len(block_out_channel_mults)
        )
    elif len(block_out_channel_mults) > len(down_block_types):
        block_out_channel_mults = block_out_channel_mults[: len(down_block_types)]
    n_mults = len(block_out_channel_mults)
    self.channel_mults = prod(block_out_channel_mults)
    self.dimensions = dimensions
    self.in_channels = in_channels
    self.n_channels = n_channels
    self.latent_channels = out_channels
    if isinstance(num_layers_per_block, int):
        num_layers_per_block = (num_layers_per_block,) * n_mults
    else:
        num_layers_per_block = tuple(num_layers_per_block)
        if len(num_layers_per_block) < len(down_block_types):
            num_layers_per_block += (num_layers_per_block[-1],) * (
                len(down_block_types) - len(num_layers_per_block)
            )

    # One sampler entry per level; the last governs channel bookkeeping only
    downsample_types = broadcast(
        downsample_type,
        n_mults,
        "downsample_type",
        None,
        f"[{n_mults} level(s)]",
    )
    downsample_clss = [
        require_cls(name, DOWNSAMPLE_FUNCTIONS, "downsampling type")
        for name in downsample_types
    ]

    # Block positions in order of data flow
    n_pos = n_mults + len(mid_block_types)
    path = f"[{n_mults} level(s), {len(mid_block_types)} mid block(s)]"
    has_attention = [
        block_type in ATTN_BLOCK_MAP
        for block_type in (*down_block_types, *mid_block_types)
    ]
    res_args = broadcast_kwargs(res_args, n_pos, context=path)
    attn_args = broadcast_kwargs(
        attn_args,
        n_pos,
        mask=has_attention,
        opaque=INTRA_BLOCK_ATTN_ARGS,
        context=path,
    )

    self.conv_in = DIM_TO_CONV_MAP[dimensions](
        in_channels,
        n_channels,
        kernel_size=kernel_size,
        stride=1,
        padding="same",
    )

    self.down_blocks = nn.ModuleList()
    ins = n_channels
    for i in range(n_mults):
        outs = ins
        if not downsample_clss[i].changes_channels:
            outs = int(ins * block_out_channel_mults[i])
        stage = nn.Sequential()
        for _ in range(num_layers_per_block[i]):
            down_block = BLOCK_MAP[down_block_types[i]](
                dimensions=dimensions,
                in_channels=ins,
                out_channels=outs,
                res_args=res_args[i],
                attn_args=attn_args[i],
            )
            stage.append(down_block)
            ins = outs
        self.down_blocks.append(stage)

        if i < n_mults - 1:
            if not downsample_clss[i].changes_channels:
                self.down_blocks.append(downsample_clss[i](dimensions, ins))
            else:
                self.down_blocks.append(
                    downsample_clss[i](
                        dimensions,
                        ins,
                        outs := int(ins * block_out_channel_mults[i]),
                    )
                )
                ins = outs
    self.bottleneck_channels = outs

    self.mid_blocks = nn.ModuleList([])
    for j, mid_block_type in enumerate(mid_block_types):
        mid_block = BLOCK_MAP[mid_block_type](
            dimensions=dimensions,
            channels=outs,
            res_args=res_args[n_mults + j],
            attn_args=attn_args[n_mults + j],
        )
        self.mid_blocks.append(mid_block)

    self.levels = (len(self.down_blocks) + 1) // 2
    self.out_channels = out_channels
    self.out_block = BLOCK_MAP[out_block_type](
        dimensions=dimensions,
        in_channels=outs,
        out_channels=self.out_channels,
        act_fn=act_fn,
        norm_type=norm_type,
        num_groups=num_groups,
        kernel_size=kernel_size,
        stride=1,
        padding="same",
    )
    self.out_shortcut = out_shortcut
    self.shortcut_groups = outs // self.out_channels

f property

f: int

Compression factor of the encoder.

Raises:

Type	Description
ValueError	If a sampling block does not scale by a constant factor.

forward

forward(x)

Forward pass.

Source code in src/chuchichaestli/models/autoencoder/encoder.py

def forward(self, x):
    """Forward pass."""
    h = self.conv_in(x)
    for block in self.down_blocks:
        h = block(h)
    for block in self.mid_blocks:
        h = block(h)
    if self.out_shortcut:
        shortcut = h.unflatten(1, (-1, self.shortcut_groups)).mean(dim=2)
        h = self.out_block(h) + shortcut
    else:
        h = self.out_block(h)
    return h

EncoderLike

Bases: Protocol

Interface an autoencoder expects from its encoding component.

Any module exposing these members can be passed to Autoencoder; inheriting from Encoder is not required. Attributes absent from a custom component simply skip the consistency checks that read them.

Attributes:

Name	Type	Description
dimensions	int	Number of spatial dimensions.
in_channels	int	Number of input channels.
n_channels	int	Number of channels after the input convolution.
out_channels	int	Number of channels emitted.
latent_channels	int	Number of latent channels out_channels stands for; smaller than out_channels when several channels describe one latent channel (a mean and a variance, say).
bottleneck_channels	int	Number of channels entering the output block.
channel_mults	int	Total channel multiplication across all levels.
levels	int	Number of spatially hierarchical levels.

Methods:

Name	Description
__call__	Encode an input tensor.
buffers	Iterate over the component's buffers.
parameters	Iterate over the component's parameters.

f property

f: int

Spatial compression factor.

__call__

__call__(x: torch.Tensor) -> torch.Tensor

Encode an input tensor.

Parameters:

Name	Type	Description	Default
x	torch.Tensor	Input tensor.	required

Source code in src/chuchichaestli/models/autoencoder/traits.py

def __call__(self, x: torch.Tensor) -> torch.Tensor:
    """Encode an input tensor.

    Args:
        x: Input tensor.
    """
    ...

buffers

buffers() -> Iterator[torch.Tensor]

Iterate over the component's buffers.

Source code in src/chuchichaestli/models/autoencoder/traits.py

def buffers(self) -> Iterator[torch.Tensor]:
    """Iterate over the component's buffers."""
    ...

parameters

parameters() -> Iterator[nn.Parameter]

Iterate over the component's parameters.

Source code in src/chuchichaestli/models/autoencoder/traits.py

def parameters(self) -> Iterator[nn.Parameter]:
    """Iterate over the component's parameters."""
    ...

LiteVAE

LiteVAE(
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
)

Bases: VAE

Lightweight wavelet-domain variational autoencoder.

Replaces the convolutional encoder of a variational autoencoder with a multi-level wavelet transform followed by small per-scale feature extractors and an aggregation network, which reaches the same reconstruction quality with a fraction of the encoder parameters.

Attributes:

Name	Type	Description
encoder_cls	type	Encoder class that build instantiates.
decoder_cls	type	Decoder class that build instantiates.

Assemble a wavelet-domain variational autoencoder from its components.

Parameters:

Name	Type	Description	Default
encoder	EncoderLike	Encoding component; must double its latent channels.	required
decoder	DecoderLike	Decoding component, expanding the latent space to the output.	required
latent_proj	nn.Module | bool	Projection between encoder and latent space; omitted by default, which reconstructs better than a pointwise convolution.	False
latent_deproj	nn.Module | bool	Projection between latent space and decoder.	False

Methods:

Name	Description
build	Build a model from architecture arguments, components included.

Source code in src/chuchichaestli/models/autoencoder/litevae.py

def __init__(
    self,
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
):
    """Assemble a wavelet-domain variational autoencoder from its components.

    Args:
        encoder: Encoding component; must double its latent channels.
        decoder: Decoding component, expanding the latent space to the output.
        latent_proj: Projection between encoder and latent space; omitted by
            default, which reconstructs better than a pointwise convolution.
        latent_deproj: Projection between latent space and decoder.
    """
    super().__init__(encoder, decoder, latent_proj, latent_deproj)

build classmethod

build(
    dimensions: int = 2,
    in_channels: int = 3,
    out_channels: int = 3,
    latent_dim: int = 12,
    decoder_n_channels: int = 512,
    **kwargs,
) -> LiteVAE

Build a model from architecture arguments, components included.

Every architecture argument is open; the published sizes are available as the LiteVAE_S, LiteVAE_B, LiteVAE_M and LiteVAE_L subclasses.

Parameters:

Name	Type	Description	Default
dimensions	int	Number of dimensions for the model.	2
in_channels	int	Number of input channels.	3
out_channels	int	Number of output channels.	3
latent_dim	int	Number of channels in the latent space.	12
decoder_n_channels	int	Number of channels the decoder starts from.	512
kwargs		Further arguments for Autoencoder.build.	{}

Source code in src/chuchichaestli/models/autoencoder/litevae.py

@classmethod
def build(
    cls,
    dimensions: int = 2,
    in_channels: int = 3,
    out_channels: int = 3,
    latent_dim: int = 12,
    decoder_n_channels: int = 512,
    **kwargs,
) -> "LiteVAE":
    """Build a model from architecture arguments, components included.

    Every architecture argument is open; the published sizes are available
    as the `LiteVAE_S`, `LiteVAE_B`, `LiteVAE_M` and `LiteVAE_L` subclasses.

    Args:
        dimensions: Number of dimensions for the model.
        in_channels: Number of input channels.
        out_channels: Number of output channels.
        latent_dim: Number of channels in the latent space.
        decoder_n_channels: Number of channels the decoder starts from.
        kwargs: Further arguments for `Autoencoder.build`.
    """
    return super().build(
        dimensions=dimensions,
        in_channels=in_channels,
        out_channels=out_channels,
        latent_dim=latent_dim,
        decoder_n_channels=decoder_n_channels,
        **kwargs,
    )

LiteVAEDecoder

LiteVAEDecoder(
    dimensions: int = 2,
    in_channels: int = 12,
    n_channels: int = 512,
    out_channels: int = 3,
    up_block_types: Sequence[str] = (
        "SMConvAutoencoderUpBlock",
    )
    * 4,
    mid_block_types: Sequence[str] = (
        "SMConvAutoencoderMidBlock",
        "AttnAutoencoderMidBlock",
    ),
    block_out_channel_mults: Sequence[int] = (1, 1, 2, 2),
    num_layers_per_block: int | Sequence[int] = 3,
    act_fn: ActivationTypes = "silu",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    **kwargs,
)

Bases: Decoder

Decoding component of a LiteVAE.

Structurally the decoder of a standard latent diffusion autoencoder, with self-modulated convolutions in place of the normalizations of its residual blocks; see Decoder for the architecture arguments. The attention block of the bottleneck keeps its own normalization.

Constructor.

Parameters:

Name	Type	Description	Default
dimensions	int	Number of dimensions.	2
in_channels	int	Number of latent channels consumed.	12
n_channels	int	Number of channels the decoder starts from.	512
out_channels	int	Number of output channels.	3
up_block_types	Sequence[str]	Type of up blocks to use for each level.	('SMConvAutoencoderUpBlock',) * 4
mid_block_types	Sequence[str]	Type of blocks to use after the input.	('SMConvAutoencoderMidBlock', 'AttnAutoencoderMidBlock')
block_out_channel_mults	Sequence[int]	Divisor for the output channels of each block.	(1, 1, 2, 2)
num_layers_per_block	int | Sequence[int]	Number of blocks per level.	3
act_fn	ActivationTypes	Activation function.	'silu'
num_groups	int	Number of groups for the normalization; also the default for the residual blocks, which are narrower here than usual.	8
kernel_size	int	Kernel size for the output convolution.	3
res_args	dict	Arguments for the residual blocks.	{}
kwargs		Further architecture arguments for Decoder.	{}

Source code in src/chuchichaestli/models/autoencoder/litevae.py

def __init__(
    self,
    dimensions: int = 2,
    in_channels: int = 12,
    n_channels: int = 512,
    out_channels: int = 3,
    up_block_types: Sequence[str] = ("SMConvAutoencoderUpBlock",) * 4,
    mid_block_types: Sequence[str] = (
        "SMConvAutoencoderMidBlock",
        "AttnAutoencoderMidBlock",
    ),
    block_out_channel_mults: Sequence[int] = (1, 1, 2, 2),
    num_layers_per_block: int | Sequence[int] = 3,
    act_fn: ActivationTypes = "silu",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    **kwargs,
):
    """Constructor.

    Args:
        dimensions: Number of dimensions.
        in_channels: Number of latent channels consumed.
        n_channels: Number of channels the decoder starts from.
        out_channels: Number of output channels.
        up_block_types: Type of up blocks to use for each level.
        mid_block_types: Type of blocks to use after the input.
        block_out_channel_mults: Divisor for the output channels of each block.
        num_layers_per_block: Number of blocks per level.
        act_fn: Activation function.
        num_groups: Number of groups for the normalization; also the default
            for the residual blocks, which are narrower here than usual.
        kernel_size: Kernel size for the output convolution.
        res_args: Arguments for the residual blocks.
        kwargs: Further architecture arguments for `Decoder`.
    """
    super().__init__(
        dimensions=dimensions,
        in_channels=in_channels,
        n_channels=n_channels,
        out_channels=out_channels,
        up_block_types=up_block_types,
        mid_block_types=mid_block_types,
        block_out_channel_mults=block_out_channel_mults,
        num_layers_per_block=num_layers_per_block,
        act_fn=act_fn,
        num_groups=num_groups,
        kernel_size=kernel_size,
        res_args={"res_groups": num_groups, **res_args},
        **kwargs,
    )
    # the output block carries a normalization too, so it is replaced in kind
    self.out_block = SMConvBlock(
        dimensions,
        self.out_block.conv.in_channels,
        out_channels,
        act_fn=act_fn,
        kernel_size=kernel_size,
        stride=1,
        padding="same",
    )

LiteVAEEncoder

LiteVAEEncoder(
    dimensions: int = 2,
    in_channels: int = 3,
    n_channels: int = 32,
    out_channels: int = 12,
    dwt_levels: int = 3,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "periodization",
    extractor_channel_mults: Sequence[int] = (3, 1),
    extractor_num_blocks: int = 4,
    aggregator_channels: int | None = None,
    aggregator_channel_mults: Sequence[int] = (3, 1),
    aggregator_num_blocks: int = 4,
    act_fn: ActivationTypes = "silu",
    norm_type: NormTypes = "group",
    num_groups: int = 8,
    res_args: dict = {},
    attn_args: dict = {},
    extractor_args: dict = {},
    aggregator_args: dict = {},
)

Bases: nn.Module

Wavelet-domain encoding component of a LiteVAE.

The input is decomposed by a multi-level wavelet transform, every level is processed by its own constant-resolution U-Net, the results are pooled to a common resolution and a second U-Net aggregates them into the latent code. All spatial compression comes from the wavelet transform, so the U-Nets keep their resolution and only widen between levels.

Satisfies EncoderLike, so it can be passed to VAE and Autoencoder without inheriting from Encoder. Like VAEEncoder it emits twice its latent channels, a mean and a variance.

Constructor.

Parameters:

Name	Type	Description	Default
dimensions	int	Number of dimensions.	2
in_channels	int	Number of input channels.	3
n_channels	int	Number of channels in the feature extractors.	32
out_channels	int	Number of latent channels; twice as many are emitted.	12
dwt_levels	int	Number of wavelet levels, and hence the compression factor.	3
wavelet	str | Wavelet	Wavelet to decompose with, by name or as a Wavelet.	'haar'
mode	ExtensionModeTypes	Signal extension mode; the default is critically sampled, and hence the only one that halves each axis exactly.	'periodization'
extractor_channel_mults	Sequence[int]	Channel multiplier per level of an extractor, as the ratio to the level before it.	(3, 1)
extractor_num_blocks	int	Number of blocks per level of an extractor.	4
aggregator_channels	int | None	Number of channels in the aggregator; as many as the extractors if omitted.	None
aggregator_channel_mults	Sequence[int]	Channel multiplier per level of the aggregator, as the ratio to the level before it.	(3, 1)
aggregator_num_blocks	int	Number of blocks per level of the aggregator.	4
act_fn	ActivationTypes	Activation function.	'silu'
norm_type	NormTypes	Normalization type.	'group'
num_groups	int	Number of groups for the normalization.	8
res_args	dict	Arguments for the residual blocks; those a UNet does not take are ignored.	{}
attn_args	dict	Arguments for the attention blocks; those a UNet does not take are ignored.	{}
extractor_args	dict	Further architecture arguments for the extractors.	{}
aggregator_args	dict	Further architecture arguments for the aggregator.	{}

Raises:

Type	Description
ValueError	If fewer than one wavelet level is requested, or if the extension mode is not critically sampled.

Methods:

Name	Description
forward	Forward pass through the encoder.

Attributes:

Name	Type	Description
f	int	Compression factor of the encoder, which the wavelet levels alone set.

Source code in src/chuchichaestli/models/autoencoder/litevae.py

def __init__(
    self,
    dimensions: int = 2,
    in_channels: int = 3,
    n_channels: int = 32,
    out_channels: int = 12,
    dwt_levels: int = 3,
    wavelet: str | Wavelet = "haar",
    mode: ExtensionModeTypes = "periodization",
    extractor_channel_mults: Sequence[int] = (3, 1),
    extractor_num_blocks: int = 4,
    aggregator_channels: int | None = None,
    aggregator_channel_mults: Sequence[int] = (3, 1),
    aggregator_num_blocks: int = 4,
    act_fn: ActivationTypes = "silu",
    norm_type: NormTypes = "group",
    num_groups: int = 8,
    res_args: dict = {},
    attn_args: dict = {},
    extractor_args: dict = {},
    aggregator_args: dict = {},
):
    """Constructor.

    Args:
        dimensions: Number of dimensions.
        in_channels: Number of input channels.
        n_channels: Number of channels in the feature extractors.
        out_channels: Number of latent channels; twice as many are emitted.
        dwt_levels: Number of wavelet levels, and hence the compression factor.
        wavelet: Wavelet to decompose with, by name or as a `Wavelet`.
        mode: Signal extension mode; the default is critically sampled, and
            hence the only one that halves each axis exactly.
        extractor_channel_mults: Channel multiplier per level of an extractor,
            as the ratio to the level before it.
        extractor_num_blocks: Number of blocks per level of an extractor.
        aggregator_channels: Number of channels in the aggregator; as many as
            the extractors if omitted.
        aggregator_channel_mults: Channel multiplier per level of the
            aggregator, as the ratio to the level before it.
        aggregator_num_blocks: Number of blocks per level of the aggregator.
        act_fn: Activation function.
        norm_type: Normalization type.
        num_groups: Number of groups for the normalization.
        res_args: Arguments for the residual blocks; those a `UNet` does
            not take are ignored.
        attn_args: Arguments for the attention blocks; those a `UNet` does
            not take are ignored.
        extractor_args: Further architecture arguments for the extractors.
        aggregator_args: Further architecture arguments for the aggregator.

    Raises:
        ValueError: If fewer than one wavelet level is requested, or if the
            extension mode is not critically sampled.
    """
    super().__init__()
    if dwt_levels < 1:
        raise ValueError(f"A wavelet encoder needs at least one level; got {dwt_levels}.")
    if mode != "periodization":
        raise ValueError(
            f"A wavelet encoder pools every level to the resolution of the"
            f" coarsest, which only a critically sampled transform reaches;"
            f" got {mode!r} instead of 'periodization'."
        )
    aggregator_channels = (
        n_channels if aggregator_channels is None else aggregator_channels
    )
    band_channels = in_channels * 2**dimensions

    self.dimensions = dimensions
    self.in_channels = in_channels
    self.n_channels = n_channels
    self.latent_channels = out_channels
    self.out_channels = 2 * out_channels
    self.levels = dwt_levels
    self.channel_mults = prod(aggregator_channel_mults)
    self.bottleneck_channels = aggregator_channels * self.channel_mults

    self.dwt = MultilevelWaveletTransformND(
        dimensions, wavelet, mode, "subband", levels=dwt_levels
    )
    shared = {
        "act_fn": act_fn,
        "norm_type": norm_type,
        "groups": num_groups,
        "res_groups": num_groups,
        "downsample_type": "ChannelResample",
        "upsample_type": "ChannelResample",
        "time_embedding": None,
        **res_args,
        **{
            f"attn_{name}": attn_args[name]
            for name in (
                    "n_heads",
                    "head_dim",
                    "dropout_p",
                    "norm_type",
                    "groups",
                    "kernel_size",
            )  # UNet attn_args
            if name in attn_args
        },
    }
    extractor_levels = len(extractor_channel_mults)
    aggregator_levels = len(aggregator_channel_mults)

    self.extractors = nn.ModuleList(
        UNet(
            dimensions=dimensions,
            in_channels=band_channels,
            n_channels=n_channels,
            out_channels=band_channels,
            down_block_types=("DownBlock",) * extractor_levels,
            mid_block_type="MidBlock",
            up_block_types=("UpBlock",) * extractor_levels,
            block_out_channel_mults=tuple(extractor_channel_mults),
            num_blocks_per_level=extractor_num_blocks,
            **{**shared, **extractor_args},
        )
        for _ in range(dwt_levels)
    )
    # every level is pooled down to the resolution of the coarsest one
    self.pools = nn.ModuleList(
        AvgPool(dimensions, kernel_size=factor, stride=factor, padding=0)
        if (factor := 2 ** (dwt_levels - 1 - level)) > 1
        else nn.Identity()
        for level in range(dwt_levels)
    )
    self.aggregator = UNet(
        dimensions=dimensions,
        in_channels=dwt_levels * band_channels,
        n_channels=aggregator_channels,
        out_channels=self.out_channels,
        down_block_types=("DownBlock",) * aggregator_levels,
        mid_block_type="MidBlock",
        up_block_types=("UpBlock",) * aggregator_levels,
        block_out_channel_mults=tuple(aggregator_channel_mults),
        num_blocks_per_level=aggregator_num_blocks,
        **{**shared, **aggregator_args},
    )

f property

f: int

Compression factor of the encoder, which the wavelet levels alone set.

forward

forward(x: torch.Tensor) -> torch.Tensor

Forward pass through the encoder.

Source code in src/chuchichaestli/models/autoencoder/litevae.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass through the encoder."""
    bands = self.dwt(x)
    features = [
        pool(extractor(band))
        for extractor, pool, band in zip(
            self.extractors, self.pools, bands, strict=True
        )
    ]
    return self.aggregator(torch.cat(features, dim=1))

LiteVAE_B

LiteVAE_B(
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
)

Bases: LiteVAE

LiteVAE with the base published encoder, about 7M parameters (default).

Source code in src/chuchichaestli/models/autoencoder/litevae.py

def __init__(
    self,
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
):
    """Assemble a wavelet-domain variational autoencoder from its components.

    Args:
        encoder: Encoding component; must double its latent channels.
        decoder: Decoding component, expanding the latent space to the output.
        latent_proj: Projection between encoder and latent space; omitted by
            default, which reconstructs better than a pointwise convolution.
        latent_deproj: Projection between latent space and decoder.
    """
    super().__init__(encoder, decoder, latent_proj, latent_deproj)

LiteVAE_L

LiteVAE_L(
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
)

Bases: LiteVAE

LiteVAE with the large published encoder, about 42M parameters.

Source code in src/chuchichaestli/models/autoencoder/litevae.py

def __init__(
    self,
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
):
    """Assemble a wavelet-domain variational autoencoder from its components.

    Args:
        encoder: Encoding component; must double its latent channels.
        decoder: Decoding component, expanding the latent space to the output.
        latent_proj: Projection between encoder and latent space; omitted by
            default, which reconstructs better than a pointwise convolution.
        latent_deproj: Projection between latent space and decoder.
    """
    super().__init__(encoder, decoder, latent_proj, latent_deproj)

LiteVAE_M

LiteVAE_M(
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
)

Bases: LiteVAE

LiteVAE with the medium published encoder, about 32M parameters.

Source code in src/chuchichaestli/models/autoencoder/litevae.py

def __init__(
    self,
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
):
    """Assemble a wavelet-domain variational autoencoder from its components.

    Args:
        encoder: Encoding component; must double its latent channels.
        decoder: Decoding component, expanding the latent space to the output.
        latent_proj: Projection between encoder and latent space; omitted by
            default, which reconstructs better than a pointwise convolution.
        latent_deproj: Projection between latent space and decoder.
    """
    super().__init__(encoder, decoder, latent_proj, latent_deproj)

LiteVAE_S

LiteVAE_S(
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
)

Bases: LiteVAE

LiteVAE with the small published encoder, about 1M parameters.

Source code in src/chuchichaestli/models/autoencoder/litevae.py

def __init__(
    self,
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = False,
    latent_deproj: nn.Module | bool = False,
):
    """Assemble a wavelet-domain variational autoencoder from its components.

    Args:
        encoder: Encoding component; must double its latent channels.
        decoder: Decoding component, expanding the latent space to the output.
        latent_proj: Projection between encoder and latent space; omitted by
            default, which reconstructs better than a pointwise convolution.
        latent_deproj: Projection between latent space and decoder.
    """
    super().__init__(encoder, decoder, latent_proj, latent_deproj)

VAE

VAE(
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = True,
    latent_deproj: nn.Module | bool = True,
)

Bases: Autoencoder

Flexible variational autoencoder implementation.

The architecture consists of an encoder-decoder structure. The encoder chains several residual and downsampling blocks. Each downsampling block separates the encoder into spatially hierarchical levels. The encoder ends in bottleneck blocks (optionally including attention blocks and a convolutional layer) and projects the input into latent space. Latent mean and variance are sampled from a Gaussian and passed to the decoder. The decoder is built with residual convolutional and upsampling blocks, and expands from the latent space to the image domain.

Attributes:

Name	Type	Description
encoder_cls	type	Encoder class that build instantiates.
decoder_cls	type	Decoder class that build instantiates.

Assemble a variational autoencoder from its two components.

Parameters:

Name	Type	Description	Default
encoder	EncoderLike	Encoding component; must double its latent channels.	required
decoder	DecoderLike	Decoding component, expanding the latent space to the output.	required
latent_proj	nn.Module | bool	Projection between encoder and latent space (see Autoencoder).	True
latent_deproj	nn.Module | bool	Projection between latent space and decoder.	True

Raises:

Type	Description
ValueError	If the encoder does not double its latent channels.

Methods:

Name	Description
encode	Encode the input.
forward	Forward pass through the model.
kl_divergence	Compute the KL divergence between posterior and a multivariate Gaussian.

Source code in src/chuchichaestli/models/autoencoder/vae.py

def __init__(
    self,
    encoder: EncoderLike,
    decoder: DecoderLike,
    latent_proj: nn.Module | bool = True,
    latent_deproj: nn.Module | bool = True,
):
    """Assemble a variational autoencoder from its two components.

    Args:
        encoder: Encoding component; must double its latent channels.
        decoder: Decoding component, expanding the latent space to the output.
        latent_proj: Projection between encoder and latent space (see `Autoencoder`).
        latent_deproj: Projection between latent space and decoder.

    Raises:
        ValueError: If the encoder does not double its latent channels.
    """
    latent = getattr(encoder, "latent_channels", encoder.out_channels)
    if encoder.out_channels != 2 * latent:
        raise ValueError(
            f"VAE reads a mean and a variance from the latent channels, so the"
            f" encoder must emit twice its {latent} latent channels; it emits"
            f" {encoder.out_channels} (e.g. use a VAEEncoder)."
        )
    super().__init__(encoder, decoder, latent_proj, latent_deproj)
    self.softplus = nn.Softplus()

encode

encode(
    x: torch.Tensor, eps: float = 1e-12
) -> MultivariateNormal

Encode the input.

Parameters:

Name	Type	Description	Default
x	torch.Tensor	Input tensor.	required
eps	float	Small constant value for numerical stability.	1e-12

Returns:

Type	Description
MultivariateNormal	Multivariate normal posterior distribution

Source code in src/chuchichaestli/models/autoencoder/vae.py

def encode(self, x: torch.Tensor, eps: float = 1e-12) -> MultivariateNormal:
    """Encode the input.

    Args:
        x: Input tensor.
        eps: Small constant value for numerical stability.

    Returns:
        Multivariate normal posterior distribution
    """
    z = self.encoder(x)
    z = self.latent_proj(z) if self.latent_proj is not None else z
    mean, log_var = z.chunk(2, dim=1)
    scale = self.softplus(log_var) + eps
    scale_tril = torch.diag_embed(scale)
    return MultivariateNormal(mean, scale_tril=scale_tril)

forward

forward(
    x: torch.Tensor,
    sample_posterior: bool = True,
    eps: float = 1e-12,
) -> tuple[torch.Tensor, torch.Tensor]

Forward pass through the model.

Source code in src/chuchichaestli/models/autoencoder/vae.py

def forward(
    self, x: torch.Tensor, sample_posterior: bool = True, eps: float = 1e-12
) -> tuple[torch.Tensor, torch.Tensor]:
    """Forward pass through the model."""
    posterior = self.encode(x, eps=eps)
    if sample_posterior:
        z = posterior.rsample()
    else:
        z = posterior.mode
    return self.decode(z), posterior

kl_divergence staticmethod

kl_divergence(
    posterior: torch.distributions.MultivariateNormal,
)

Compute the KL divergence between posterior and a multivariate Gaussian.

Source code in src/chuchichaestli/models/autoencoder/vae.py

@staticmethod
def kl_divergence(posterior: torch.distributions.MultivariateNormal):
    """Compute the KL divergence between posterior and a multivariate Gaussian."""
    device = posterior.mean.device
    dtype = posterior.mean.dtype
    zeros = torch.zeros_like(posterior.mean)
    eye = torch.eye(posterior.mean.shape[-1], device=device, dtype=dtype)
    return kl.kl_divergence(posterior, MultivariateNormal(zeros, eye))

VAEDecoder

VAEDecoder(
    dimensions: int = 2,
    in_channels: int = 1,
    n_channels: int = 512,
    out_channels: int = 1,
    in_block_type: DecoderInBlockTypes = "DecoderInBlock",
    mid_block_types: Sequence[AutoencoderMidBlockTypes] = (
        "AutoencoderMidBlock",
        "AttnAutoencoderMidBlock",
    ),
    up_block_types: Sequence[AutoencoderUpBlockTypes] = (
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
    ),
    block_out_channel_mults: Sequence[int] = (1, 2, 2, 2),
    num_layers_per_block: int | Sequence[int] = 3,
    upsample_type: UpsampleTypes
    | Sequence[UpsampleTypes] = "UpsampleInterpolate",
    act_fn: ActivationTypes = "silu",
    norm_type: NormTypes = "group",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    attn_args: dict = {},
    in_shortcut: bool = False,
)

Bases: Decoder

Decoding component of a variational autoencoder.

Consumes a latent code that sampling has already collapsed to a single set of channels, so structurally it is a plain Decoder; see there for the architecture arguments.

Source code in src/chuchichaestli/models/autoencoder/decoder.py

def __init__(
    self,
    dimensions: int = 2,
    in_channels: int = 1,
    n_channels: int = 512,
    out_channels: int = 1,
    in_block_type: DecoderInBlockTypes = "DecoderInBlock",
    mid_block_types: Sequence[AutoencoderMidBlockTypes] = (
        "AutoencoderMidBlock",
        "AttnAutoencoderMidBlock",
    ),
    up_block_types: Sequence[AutoencoderUpBlockTypes] = (
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
        "AutoencoderUpBlock",
    ),
    block_out_channel_mults: Sequence[int] = (1, 2, 2, 2),
    num_layers_per_block: int | Sequence[int] = 3,
    upsample_type: UpsampleTypes | Sequence[UpsampleTypes] = "UpsampleInterpolate",
    act_fn: ActivationTypes = "silu",
    norm_type: NormTypes = "group",
    num_groups: int = 8,
    kernel_size: int = 3,
    res_args: dict = {},
    attn_args: dict = {},
    in_shortcut: bool = False,
):
    """Decoder implementation.

    Args:
        dimensions: Number of dimensions.
        in_channels: Number of input channels (latent space).
        n_channels: Number of channels for first block.
        out_channels: Number of output channels.
        in_block_type: Type of block for output (latent space).
        mid_block_types: Type of blocks to use before the output.
        up_block_types: Type of up blocks to use for each level.
        block_out_channel_mults: Multiplier for output channels of each block.
        num_layers_per_block: Number of blocks per level (blocks are repeated if `>1`).
        upsample_type: Type of upsampling block, per level (see
            `chuchichaestli.models.upsampling` for details).
        act_fn: Activation function for the output layers
            (see `chuchichaestli.models.activations` for details).
        norm_type: Normalization type for the output layer.
        num_groups: Number of groups for normalization in the output layer.
        kernel_size: Kernel size for the output convolution.
        res_args: Arguments for residual blocks. Each value is either a
            single value or one per block position, in order of data flow.
        attn_args: Arguments for attention blocks, per block position or per
            block that has attention. `norm_type`, `scales`, `context_args`
            and `local_args` are passed through unchanged.
        in_shortcut: Whether to use a shortcut for the input block.
    """
    super().__init__()

    block_out_channel_mults = tuple(block_out_channel_mults)
    if len(block_out_channel_mults) < len(up_block_types):
        block_out_channel_mults += (1,) * (
            len(up_block_types) - len(block_out_channel_mults)
        )
    elif len(block_out_channel_mults) > len(up_block_types):
        block_out_channel_mults = block_out_channel_mults[: len(up_block_types)]
    n_mults = len(block_out_channel_mults)
    self.channel_mults = prod(block_out_channel_mults)
    self.dimensions = dimensions
    self.in_channels = in_channels
    self.n_channels = n_channels
    self.out_channels = out_channels
    if isinstance(num_layers_per_block, int):
        num_layers_per_block = (num_layers_per_block,) * n_mults
    else:
        num_layers_per_block = tuple(num_layers_per_block)
        if len(num_layers_per_block) < len(up_block_types):
            num_layers_per_block += (num_layers_per_block[-1],) * (
                len(up_block_types) - len(num_layers_per_block)
            )

    # One sampler entry per level; the last governs channel bookkeeping only
    upsample_types = broadcast(
        upsample_type, n_mults, "upsample_type", None, f"[{n_mults} level(s)]"
    )
    upsample_clss = [
        require_cls(name, UPSAMPLE_FUNCTIONS, "upsampling type")
        for name in upsample_types
    ]

    # Block positions in order of data flow
    n_pos = n_mults + len(mid_block_types)
    path = f"[{len(mid_block_types)} mid block(s), {n_mults} level(s)]"
    has_attention = [
        block_type in ATTN_BLOCK_MAP
        for block_type in (*mid_block_types, *up_block_types)
    ]
    res_args = broadcast_kwargs(res_args, n_pos, context=path)
    attn_args = broadcast_kwargs(
        attn_args,
        n_pos,
        mask=has_attention,
        opaque=INTRA_BLOCK_ATTN_ARGS,
        context=path,
    )

    self.in_block = BLOCK_MAP[in_block_type](
        dimensions=dimensions,
        in_channels=in_channels,
        out_channels=n_channels,
        kernel_size=kernel_size,
    )

    self.mid_blocks = nn.ModuleList([])
    for j, mid_block_type in enumerate(mid_block_types):
        mid_block = BLOCK_MAP[mid_block_type](
            dimensions=dimensions,
            channels=n_channels,
            res_args=res_args[j],
            attn_args=attn_args[j],
        )
        self.mid_blocks.append(mid_block)

    self.up_blocks = nn.ModuleList([])
    ins = n_channels
    for i in range(n_mults):
        outs = ins
        if not upsample_clss[i].changes_channels:
            outs = int(ins // block_out_channel_mults[i])
        stage = nn.Sequential()
        for _ in range(num_layers_per_block[i]):
            up_block = BLOCK_MAP[up_block_types[i]](
                dimensions=dimensions,
                in_channels=ins,
                out_channels=outs,
                res_args=res_args[len(mid_block_types) + i],
                attn_args=attn_args[len(mid_block_types) + i],
            )
            stage.append(up_block)
            ins = outs
        self.up_blocks.append(stage)

        if i < n_mults - 1:
            if not upsample_clss[i].changes_channels:
                self.up_blocks.append(upsample_clss[i](dimensions, outs))
            else:
                self.up_blocks.append(
                    upsample_clss[i](
                        dimensions,
                        ins,
                        outs := int(ins // block_out_channel_mults[i]),
                    )
                )
                ins = outs
    self.levels = (len(self.up_blocks) + 1) // 2
    self.out_block = NormActConvBlock(
        dimensions=dimensions,
        in_channels=outs,
        out_channels=out_channels,
        act_fn=act_fn,
        norm_type=norm_type,
        num_groups=num_groups,
        kernel_size=kernel_size,
        stride=1,
        padding="same",
    )
    self.in_shortcut = in_shortcut
    self.shortcut_repeats = n_channels // in_channels

VAEEncoder

VAEEncoder(
    dimensions: int = 2,
    in_channels: int = 1,
    n_channels: int = 64,
    out_channels: int = 1,
    **kwargs,
)

Bases: Encoder

Encoding component of a variational autoencoder.

Doubles the latent channels so that a mean and a variance can be read from them; see Encoder for the architecture arguments.

Constructor.

The leading arguments are named so that out_channels is doubled whether it is passed positionally or by keyword.

Parameters:

Name	Type	Description	Default
dimensions	int	Number of dimensions.	2
in_channels	int	Number of input channels.	1
n_channels	int	Number of channels in the hidden layer.	64
out_channels	int	Number of latent channels; twice as many are emitted.	1
kwargs		Further architecture arguments for Encoder.	{}

Source code in src/chuchichaestli/models/autoencoder/vae.py

def __init__(
    self,
    dimensions: int = 2,
    in_channels: int = 1,
    n_channels: int = 64,
    out_channels: int = 1,
    **kwargs,
):
    """Constructor.

    The leading arguments are named so that `out_channels` is doubled
    whether it is passed positionally or by keyword.

    Args:
        dimensions: Number of dimensions.
        in_channels: Number of input channels.
        n_channels: Number of channels in the hidden layer.
        out_channels: Number of latent channels; twice as many are emitted.
        kwargs: Further architecture arguments for `Encoder`.
    """
    super().__init__(
        dimensions=dimensions,
        in_channels=in_channels,
        n_channels=n_channels,
        out_channels=2 * out_channels,
        **kwargs,
    )
    self.latent_channels = out_channels

VQVAE

VQVAE(
    encoder: EncoderLike,
    decoder: DecoderLike,
    vq_dim: int = 64,
    vq_embeddings: int = 512,
)

Bases: Autoencoder

Flexible vector-quantized variational autoencoder implementation.

The architecture consists of an encoder-decoder structure with a codebook bottleneck. The encoder chains several residual and downsampling blocks. Each downsampling block separates the encoder into spatially hierarchical levels. The encoder ends in bottleneck blocks (optionally including attention blocks and a convolutional layer). The latent code is quantized with the codebook. The decoder is built with residual convolutional and upsampling blocks, and expands from the (quantized) latent space to the image domain.

Assemble a vector-quantized autoencoder from its two components.

Parameters:

Name	Type	Description	Default
encoder	EncoderLike	Encoding component, mapping the input to latent space.	required
decoder	DecoderLike	Decoding component, expanding the latent space to the output.	required
vq_dim	int	Size of the quantized embedding vectors.	64
vq_embeddings	int	Size of the quantization codebook.	512

Methods:

Name	Description
compute_embedding_shape	Compute the shape of the latent space.
encode	Encode the input.
forward	Forward pass through the model.

Source code in src/chuchichaestli/models/autoencoder/vqvae.py

def __init__(
    self,
    encoder: EncoderLike,
    decoder: DecoderLike,
    vq_dim: int = 64,
    vq_embeddings: int = 512,
):
    """Assemble a vector-quantized autoencoder from its two components.

    Args:
        encoder: Encoding component, mapping the input to latent space.
        decoder: Decoding component, expanding the latent space to the output.
        vq_dim: Size of the quantized embedding vectors.
        vq_embeddings: Size of the quantization codebook.
    """
    latent_dim = getattr(encoder, "latent_channels", encoder.out_channels)
    if encoder.out_channels != latent_dim:
        raise ValueError(
            f"VQVAE quantizes the latent code as it is, so the encoder must emit"
            f" its {latent_dim} latent channels; this one emits"
            f" {encoder.out_channels} (an encoder that describes each latent"
            " channel with several outputs, such as VAEEncoder, has no"
            " single code to quantize)."
        )
    super().__init__(
        encoder,
        decoder,
        latent_proj=self.projection(encoder, latent_dim, vq_dim),
        latent_deproj=self.projection(decoder, vq_dim, latent_dim),
    )
    self.vq_dim = vq_dim
    self.quantize = VectorQuantizer(
        num_embeddings=vq_embeddings, embedding_dim=vq_dim
    )

compute_embedding_shape

compute_embedding_shape(
    input_shape: tuple[int, ...], no_batch_dim: bool = False
)

Compute the shape of the latent space.

Source code in src/chuchichaestli/models/autoencoder/vqvae.py

def compute_embedding_shape(
    self, input_shape: tuple[int, ...], no_batch_dim: bool = False
):
    """Compute the shape of the latent space."""
    batch_dim = input_shape[0] if not no_batch_dim else None
    spatial_dims = tuple(dim // self.f_comp for dim in input_shape[2:])
    if batch_dim is None:
        shape = (self.vq_dim, *spatial_dims)
    else:
        shape = (batch_dim, self.vq_dim, *spatial_dims)
    return shape

encode

encode(
    x: torch.Tensor | torch.LongTensor,
    codebook_usage: bool = False,
    force_no_quant: bool = False,
    load_from_codebook: bool = False,
    shape: tuple[int, ...] | None = None,
) -> tuple[torch.Tensor, torch.Tensor, dict]

Encode the input.

Parameters:

Name	Type	Description	Default
x	torch.Tensor | torch.LongTensor	Input tensor or codebook indices (if load_from_codebook and force_no_quant).	required
codebook_usage	bool	If True, compute and return codebook usage parameters.	False
force_no_quant	bool	If True, quantization of the latent code is skipped.	False
load_from_codebook	bool	If True, indices are loaded from the codebook instead.	False
shape	tuple[int, ...] | None	Full latent shape for codebook entries (with the embedding dimension in the last axis).	None

Returns:

Type	Description
torch.Tensor	Quantized latent code
torch.Tensor	Quantization loss
dict	Codebook usage parameters (perplexity, flattened latent, codebook indices)

Source code in src/chuchichaestli/models/autoencoder/vqvae.py

def encode(
    self,
    x: torch.Tensor | torch.LongTensor,
    codebook_usage: bool = False,
    force_no_quant: bool = False,
    load_from_codebook: bool = False,
    shape: tuple[int, ...] | None = None,
) -> tuple[torch.Tensor, torch.Tensor, dict]:
    """Encode the input.

    Args:
        x: Input tensor or codebook indices (if `load_from_codebook` and `force_no_quant`).
        codebook_usage: If `True`, compute and return codebook usage parameters.
        force_no_quant: If `True`, quantization of the latent code is skipped.
        load_from_codebook: If `True`, indices are loaded from the codebook instead.
        shape: Full latent shape for codebook entries
            (with the embedding dimension in the last axis).

    Returns:
        - Quantized latent code
        - Quantization loss
        - Codebook usage parameters (perplexity, flattened latent, codebook indices)
    """
    z = self._encode(x)
    if not force_no_quant:
        z_q, loss, usage = self.quantize(z, codebook_usage=codebook_usage)
    elif load_from_codebook:
        z_q = self.quantize.get_codebook_entry(x, shape)
        loss = torch.as_tensor(0.0, device=z.device, dtype=z.dtype)
        usage = {}
    else:
        z_q = z
        loss = torch.as_tensor(0.0, device=z.device, dtype=z.dtype)
        usage = {}
    return z_q, loss, usage

forward

forward(
    x: torch.Tensor | torch.LongTensor,
    codebook_usage: bool = False,
    force_no_quant: bool = False,
    load_from_codebook: bool = False,
    shape: tuple[int, ...] | None = None,
) -> tuple[torch.Tensor, torch.Tensor, dict]

Forward pass through the model.

Parameters:

Name	Type	Description	Default
x	torch.Tensor | torch.LongTensor	Input tensor or codebook indices (if load_from_codebook and force_no_quant).	required
codebook_usage	bool	If True, compute and return codebook usage parameters.	False
force_no_quant	bool	If True, quantization of the latent code is skipped.	False
load_from_codebook	bool	If True, indices are loaded from the codebook instead.	False
shape	tuple[int, ...] | None	Full latent shape for codebook entries (with the embedding dimension in the last axis).	None

Returns:

Type	Description
torch.Tensor	Quantized latent code
torch.Tensor	Quantization loss
dict	Codebook usage parameters (perplexity, flattened latent, codebook indices)

Source code in src/chuchichaestli/models/autoencoder/vqvae.py

def forward(
    self,
    x: torch.Tensor | torch.LongTensor,
    codebook_usage: bool = False,
    force_no_quant: bool = False,
    load_from_codebook: bool = False,
    shape: tuple[int, ...] | None = None,
) -> tuple[torch.Tensor, torch.Tensor, dict]:
    """Forward pass through the model.

    Args:
        x: Input tensor or codebook indices (if `load_from_codebook` and `force_no_quant`).
        codebook_usage: If `True`, compute and return codebook usage parameters.
        force_no_quant: If `True`, quantization of the latent code is skipped.
        load_from_codebook: If `True`, indices are loaded from the codebook instead.
        shape: Full latent shape for codebook entries
            (with the embedding dimension in the last axis).

    Returns:
        - Quantized latent code
        - Quantization loss
        - Codebook usage parameters (perplexity, flattened latent, codebook indices)
    """
    z_q, loss, usage = self.encode(
        x,
        codebook_usage=codebook_usage,
        force_no_quant=force_no_quant,
        load_from_codebook=load_from_codebook,
        shape=shape,
    )
    return self.decode(z_q), loss, usage

VectorQuantizer

VectorQuantizer(
    num_embeddings: int,
    embedding_dim: int,
    beta: float = 0.25,
)

Bases: nn.Module

Vector Quantizer for VQVAE.

Initialize VectorQuantizer.

Parameters:

Name	Type	Description	Default
num_embeddings	int	Size of the codebook.	required
embedding_dim	int	Size of the embedding vectors.	required
beta	float	Commitment cost parameter for the loss.	0.25

Methods:

Name	Description
forward	Forward pass through the vector quantizer (codebook lookup).
get_codebook_entry	Get the codebook entry for the given indices.

Attributes:

Name	Type	Description
embedding_dim	int	Size of each embedding vector.
num_embeddings	int	Size of the embedding dictionary.

Source code in src/chuchichaestli/models/autoencoder/vqvae.py

def __init__(self, num_embeddings: int, embedding_dim: int, beta: float = 0.25):
    """Initialize VectorQuantizer.

    Args:
        num_embeddings: Size of the codebook.
        embedding_dim: Size of the embedding vectors.
        beta: Commitment cost parameter for the loss.
    """
    super().__init__()
    self.beta = beta
    self.embedding = nn.Embedding(num_embeddings, embedding_dim)
    self.embedding.weight.data.uniform_(-1 / num_embeddings, 1 / num_embeddings)

embedding_dim property

embedding_dim: int

Size of each embedding vector.

num_embeddings property

num_embeddings: int

Size of the embedding dictionary.

forward

forward(
    z: torch.Tensor, codebook_usage: bool = False
) -> tuple[torch.Tensor, torch.Tensor, dict]

Forward pass through the vector quantizer (codebook lookup).

Parameters:

Name	Type	Description	Default
z	torch.Tensor	Input latent codes of shape (B, embedding_dim, *S).	required
codebook_usage	bool	If True, compute and return codebook usage parameters.	False

Returns:

Type	Description
torch.Tensor	Quantized latent code
torch.Tensor	Quantization loss
dict	Codebook usage parameters (perplexity, flattened latent, codebook indices)

Source code in src/chuchichaestli/models/autoencoder/vqvae.py

def forward(
    self, z: torch.Tensor, codebook_usage: bool = False
) -> tuple[torch.Tensor, torch.Tensor, dict]:
    """Forward pass through the vector quantizer (codebook lookup).

    Args:
        z: Input latent codes of shape `(B, embedding_dim, *S)`.
        codebook_usage: If `True`, compute and return codebook usage parameters.

    Returns:
        - Quantized latent code
        - Quantization loss
        - Codebook usage parameters (perplexity, flattened latent, codebook indices)
    """
    z = z.moveaxis(1, -1).contiguous()
    z_flat = z.view(-1, self.embedding_dim)
    # quantize
    e = self.embedding.weight
    distances = (
        z_flat.pow(2).sum(dim=1, keepdim=True)
        + e.pow(2).sum(dim=1)
        - 2 * (z_flat @ e.T)
    )  # rather compute square distances (avoiding sqrt computation of torch.cdist)
    nearest_emb_idcs = torch.argmin(distances, dim=1)
    z_q = self.embedding(nearest_emb_idcs).view(z.shape)
    # loss calculation
    commitment_loss = (z_q.detach() - z).pow(2).mean()
    codebook_loss = (z_q - z.detach()).pow(2).mean()
    loss = self.beta * commitment_loss + codebook_loss
    # preserve gradient flow
    z_q = z + (z_q - z).detach()
    z_q = z_q.moveaxis(-1, 1).contiguous()
    # calculate codebook usage
    perplexity: torch.Tensor | None = None
    if codebook_usage:
        counts = torch.bincount(
            nearest_emb_idcs, minlength=self.embedding.num_embeddings
        )
        avg_probs = counts.float() / counts.sum()
        mask = avg_probs > 0
        entropy = torch.sum(avg_probs[mask] * torch.log(avg_probs[mask]))
        perplexity = torch.exp(-entropy)
    return (
        z_q,
        loss,
        {
            "perplexity": perplexity,
            "latent_flat": z_flat,
            "indices": nearest_emb_idcs,
        },
    )

get_codebook_entry

get_codebook_entry(
    indices: torch.LongTensor,
    shape: tuple[int, ...] | None = None,
) -> torch.Tensor

Get the codebook entry for the given indices.

Parameters:

Name	Type	Description	Default
indices	torch.LongTensor	Indices from the codebook.	required
shape	tuple[int, ...] | None	Full latent shape (with the embedding dimension in the last axis).	None

Source code in src/chuchichaestli/models/autoencoder/vqvae.py

def get_codebook_entry(
    self, indices: torch.LongTensor, shape: tuple[int, ...] | None = None
) -> torch.Tensor:
    """Get the codebook entry for the given indices.

    Args:
        indices: Indices from the codebook.
        shape: Full latent shape (with the embedding dimension in the last axis).
    """
    z_q = self.embedding(indices)
    if shape is not None:
        z_q = z_q.view(shape)
        z_q = z_q.moveaxis(-1, 1).contiguous()
    return z_q