residual_vector_quantizer

ResidualVectorQuantizer(
    num_levels: int,
    num_embeddings: int | Sequence[int],
    embedding_dim: int,
    beta: float = 0.25,
    **kwargs={},
)

Residual Vector Quantizer (RVQ) with straight-through estimator.

Input: […, D] (last dim = embedding_dim) Output: […, D] (sum of per-level dequantized vectors; gradients pass through x)

Metrics (logged via metrics property):

rvq_l{l}_perplexity, rvq_l{l}_usage, rvq_l{l}_bits_per_index
rvq_perplexity_mean, rvq_usage_mean, rvq_bits_per_index_sum (entropy lower bound)

Parameters

Parameters of ResidualVectorQuantizer
Name	Type	Default	Description
`num_levels`	`int`	Required	int, number of residual VQ stages (M >= 1)
`num_embeddings`	`int \| Sequence[int]`	Required	int OR sequence[int], codebook size K for each level
`embedding_dim`	`int`	Required	int, latent dimensionality D
`beta`	`float`	0.25	float, commitment coefficient per level

M = int(num_levels)

D = int(embedding_dim)

Ks

beta = float(beta)

metrics

build(input_shape)

call(
    x: keras.KerasTensor,
    return_indices: bool = False,
) -> keras.KerasTensor | tuple[keras.KerasTensor, list[keras.KerasTensor]]

Parameters

Parameters of call
Name	Type	Default	Description
`x`	`keras.KerasTensor`	Required	[..., D] latent to be quantized.
`return_indices`	`bool`	False	if True, also returns list of flat indices (one tensor per level).

Returns

Returns of call
Type	Description
`keras.KerasTensor \| tuple[keras.KerasTensor, list[keras.KerasTensor]]`	y or (y, indices_list): dequantized vector and optional per-level indices.

encode(x)

Return list of per-level flat index tensors [N] (no gradients).

decode(indices_list, original_shape)

Sum per-level code vectors from indices_list and reshape to original_shape.

get_config()