uniform_codebook_prior

uniform_codebook_prior

Fixed uniform prior over discrete codebook indices.

The VQ-VAE equivalent of GaussianPrior: samples a random codebook index uniformly and returns the corresponding embedding. No learnable parameters. The codebook is shared from the VQ posterior encoder via wire_posterior().

UniformCodebookPrior

UniformCodebookPrior(latent_dimension, num_codes, num_residual_layers, device)

Bases: PriorLatentEncoder

Fixed uniform prior over VQ codebook indices.

Samples each residual VQ layer's index uniformly from {0, ..., K-1}, then decodes to a quantized embedding via the shared codebook. No trainable parameters — the VQ equivalent of GaussianPrior.

Parameters:

Name	Type	Description	Default
latent_dimension	int	Dimension of each codebook vector.	required
num_codes	int	Number of codebook entries per layer (K).	required
num_residual_layers	int	Number of residual VQ layers.	required
device	str	Device string.	required

Source code in src/versatil/models/decoding/latent/prior/uniform_codebook_prior.py

def __init__(
    self,
    latent_dimension: int,
    num_codes: int,
    num_residual_layers: int,
    device: str,
):
    super().__init__(latent_dimension=latent_dimension, device=device)
    if latent_dimension <= 0:
        raise ValueError(
            f"latent_dimension must be positive, got {latent_dimension}."
        )
    if num_codes <= 0:
        raise ValueError(f"num_codes must be positive, got {num_codes}.")
    if num_residual_layers <= 0:
        raise ValueError(
            f"num_residual_layers must be positive, got {num_residual_layers}."
        )
    self.code_dim = latent_dimension
    self.num_codes = num_codes
    self.num_residual_layers = num_residual_layers
    self._residual_vq_reference: tuple[ResidualVQ, ...] = ()

residual_vq property

residual_vq

Return the posterior-owned VQ module without registering it as a child.

__getstate__

__getstate__()

Return copy-safe state without the posterior VQ reference.

Returns:

Type	Description
dict[str, Any]	Module state with runtime wiring cleared. The owning
dict[str, Any]	VariationalAlgorithm reconnects the copied prior to the copied
dict[str, Any]	posterior after deepcopy or unpickling.

Source code in src/versatil/models/decoding/latent/prior/uniform_codebook_prior.py

def __getstate__(self) -> dict[str, Any]:
    """Return copy-safe state without the posterior VQ reference.

    Returns:
        Module state with runtime wiring cleared. The owning
        VariationalAlgorithm reconnects the copied prior to the copied
        posterior after deepcopy or unpickling.
    """
    state = super().__getstate__()
    state["_residual_vq_reference"] = ()
    return state

get_auxiliary_output_keys

get_auxiliary_output_keys()

Uniform prior outputs only the quantized latent and sampled indices.

Source code in src/versatil/models/decoding/latent/prior/uniform_codebook_prior.py

def get_auxiliary_output_keys(self) -> set[str]:
    """Uniform prior outputs only the quantized latent and sampled indices."""
    return {
        LatentKey.PRIOR_LATENT.value,
        LatentKey.VQ_PRIOR_INDICES.value,
    }

wire_posterior

wire_posterior(posterior)

Wire shared codebook from the VQ posterior encoder.

Parameters:

Name	Type	Description	Default
posterior	PosteriorLatentEncoder	VQ posterior encoder with a residual_vq attribute.	required

Raises:

Type	Description
AttributeError	If the posterior does not expose ResidualVQ state.
ValueError	If the posterior's VQ configuration does not match.

Source code in src/versatil/models/decoding/latent/prior/uniform_codebook_prior.py

def wire_posterior(self, posterior: PosteriorLatentEncoder) -> None:
    """Wire shared codebook from the VQ posterior encoder.

    Args:
        posterior: VQ posterior encoder with a residual_vq attribute.

    Raises:
        AttributeError: If the posterior does not expose ResidualVQ state.
        ValueError: If the posterior's VQ configuration does not match.
    """
    residual_vq = getattr(posterior, "residual_vq", None)
    if residual_vq is None:
        raise AttributeError(
            f"Posterior {type(posterior).__name__} does not expose a "
            f"residual_vq attribute required by UniformCodebookPrior."
        )
    if residual_vq.code_dim != self.code_dim:
        raise ValueError(
            f"ResidualVQ code_dim ({residual_vq.code_dim}) does not match "
            f"UniformCodebookPrior code_dim ({self.code_dim})"
        )
    if residual_vq.num_codes != self.num_codes:
        raise ValueError(
            f"ResidualVQ num_codes ({residual_vq.num_codes}) does not match "
            f"UniformCodebookPrior num_codes ({self.num_codes})"
        )
    if residual_vq.num_layers != self.num_residual_layers:
        raise ValueError(
            f"ResidualVQ num_layers ({residual_vq.num_layers}) does not match "
            f"UniformCodebookPrior num_residual_layers ({self.num_residual_layers})"
        )
    self._residual_vq_reference = (residual_vq,)

forward

forward(target_latents, observations)

Sample uniform codebook indices and decode to embedding.

Parameters:

Name	Type	Description	Default
target_latents	Tensor | None	Posterior latent used only to infer batch size. None is accepted; batch size is then derived from observations.	required
observations	dict[str, Tensor]	Observation features used only to infer batch size.	required

Returns:

Type	Description
dict[str, Tensor]	Dictionary containing: - LatentKey.PRIOR_LATENT: Quantized embedding, shape (B, code_dim). - LatentKey.VQ_PRIOR_INDICES: List of per-layer indices sampled by the prior, each shape (B,). Emitted under a distinct key from the posterior's VQ_INDICES so the two do not collide when merged into the predictions dict.

Source code in src/versatil/models/decoding/latent/prior/uniform_codebook_prior.py

def forward(
    self,
    target_latents: torch.Tensor | None,
    observations: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Sample uniform codebook indices and decode to embedding.

    Args:
        target_latents: Posterior latent used only to infer batch size.
            None is accepted; batch size is then derived from observations.
        observations: Observation features used only to infer batch size.

    Returns:
        Dictionary containing:
            - LatentKey.PRIOR_LATENT: Quantized embedding, shape (B, code_dim).
            - LatentKey.VQ_PRIOR_INDICES: List of per-layer indices
                sampled by the prior, each shape (B,). Emitted under a
                distinct key from the posterior's VQ_INDICES so the two
                do not collide when merged into the predictions dict.
    """
    residual_vq = self.residual_vq
    if residual_vq is None:
        raise RuntimeError(
            "UniformCodebookPrior.residual_vq is not set. "
            "Call wire_posterior() before forward()."
        )
    if target_latents is not None:
        batch_size = target_latents.shape[0]
    else:
        batch_size = next(iter(observations.values())).shape[0]
    device = self.device

    all_indices = [
        torch.randint(0, self.num_codes, (batch_size,), device=device)  # (B,)
        for _ in range(self.num_residual_layers)
    ]

    z_q = residual_vq.decode_from_indices(all_indices)  # (B, code_dim)

    return {
        LatentKey.PRIOR_LATENT.value: z_q,
        LatentKey.VQ_PRIOR_INDICES.value: all_indices,
    }

sample_prior

sample_prior(batch_size, observations=None)

Sample latent from uniform categorical prior.

Parameters:

Name	Type	Description	Default
batch_size	int	Number of samples.	required
observations	dict[str, Tensor] | None	Unused (uniform prior is unconditional).	None

Returns:

Type	Description
Tensor	Quantized latent embedding, shape (batch_size, code_dim).

Source code in src/versatil/models/decoding/latent/prior/uniform_codebook_prior.py

def sample_prior(
    self,
    batch_size: int,
    observations: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Sample latent from uniform categorical prior.

    Args:
        batch_size: Number of samples.
        observations: Unused (uniform prior is unconditional).

    Returns:
        Quantized latent embedding, shape (batch_size, code_dim).
    """
    residual_vq = self.residual_vq
    if residual_vq is None:
        raise RuntimeError(
            "UniformCodebookPrior.residual_vq is not set. "
            "Call wire_posterior() before sample_prior()."
        )
    device = self.device
    all_indices = [
        torch.randint(0, self.num_codes, (batch_size,), device=device)  # (B,)
        for _ in range(self.num_residual_layers)
    ]
    return residual_vq.decode_from_indices(all_indices)  # (B, code_dim)