nearest_codebook

→Tensor

nearest_codebook(x: Tensor, codebook: Tensor)

source edit

Index of the closest codebook entry for each row of x.

k_i = \arg\min_j \big\| x_i - e_j \big\|_2

The search is the non-differentiable half of vector quantisation: the result is an integer field, so no gradient flows through it and none is defined. vector_quantize pairs it with the straight-through estimator to make the surrounding network trainable.

Parameters

xTensor

Query field of shape (*, D). Only the trailing axis is treated as the feature dimension; everything before it is flattened into the search's row axis.

codebookTensor

Codebook of shape (K, D).

Returns

Tensor

int64 index field of shape (*) with values in [0, K).

Notes

Distances go through lucid.cdist, whose p=2 path uses the stable expansion $\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2ab^\top$ rather than materialising an (N, K, D) difference. It does still build the (N, K) matrix, which at a large latent grid and a large codebook is the dominant allocation of a quantiser; a fused engine kernel that reduces over K without materialising it is the natural next step and would slot in behind this exact signature.

Examples

>>> import lucid
>>> import lucid.nn.functional as F
>>> codebook = lucid.tensor([[0.0, 0.0], [1.0, 1.0]])
>>> x = lucid.tensor([[0.9, 1.1], [0.1, 0.0]])
>>> F.nearest_codebook(x, codebook).tolist()
[1, 0]

Used by 2

>>> import lucid >>> import lucid.nn.functional as F >>> codebook = lucid.tensor([[0.0, 0.0], [1.0, 1.0]]) >>> x = lucid.tensor([[0.9, 1.1], [0.1, 0.0]]) >>> F.nearest_codebook(x, codebook).tolist() [1, 0]