data
VAEConfig
extends
GenerativeModelConfigVAEConfig(sample_size: int | tuple[int, int] = 32, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'silu', latent_dim: int | tuple[int, ...] = 128, down_block_channels: tuple[int, ...] = (lambda: (64, 128, 256))(), kl_weight: float = 1.0, recon_loss: Literal['mse', 'bce'] = 'mse')Configuration for every VAE / β-VAE / HVAE variant.
Parameters
latent_dimint | tuple[int, ...]= 128Bottleneck spec.
int— vanilla / β-VAE: a singlezof this dimension is projected from the flattened final encoder stage.tuple[int, ...]— hierarchical: onez_lper encoder stage. Length must equallen(down_block_channels).
down_block_channelstuple[int, ...]= (lambda: (64, 128, 256))()Channel width after each encoder downsampling
block. Length determines the number of stride-2 blocks; the
final spatial resolution is
sample_size / 2 ** len(...).kl_weightfloat= 1.0Multiplicative β in the β-VAE loss.
1.0 recovers the
vanilla ELBO; < 1 weakens the KL pull (better
reconstructions); > 1 enforces stronger disentanglement
(Higgins et al., 2017 β-VAE). In the hierarchical case, β
scales every level's KL identically.recon_lossLiteral['mse', 'bce']= 'mse'Reconstruction term —
"mse" (Gaussian likelihood,
default) or "bce" (Bernoulli likelihood for [0, 1] data).Used by 3
Constructors
1dunder
__init__
→None__init__(sample_size: int | tuple[int, int] = 32, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'silu', latent_dim: int | tuple[int, ...] = 128, down_block_channels: tuple[int, ...] = (lambda: (64, 128, 256))(), kl_weight: float = 1.0, recon_loss: Literal['mse', 'bce'] = 'mse')Properties
3True when latent_dim is a tuple — selects the HVAE topology.
Always-tuple view of latent_dim (length 1 in vanilla mode).
Sum across levels — what the decoder consumes at its top input.