VectorQuantizer

Nearest-neighbour codebook lookup with a straight-through gradient.

Maps each position of a continuous feature field to the closest entry of a learned codebook $e \in \mathbb{R}^{K \times D}$:

This is the discretisation layer of van den Oord, Vinyals, and Kavukcuoglu, "Neural Discrete Representation Learning" (2017) — the bottleneck that turns an auto-encoder into a tokeniser whose latent is an integer field a downstream model can treat as a vocabulary.

Like Linear and Embedding, the layer acts on the trailing axis: an input of shape (*, embedding_dim) produces a quantised field of the same shape and an index field of shape (*). Image models holding (N, C, H, W) should permute the channel axis last before calling and back afterwards.

Because $\arg\min$ has zero gradient almost everywhere, the forward value is routed through lucid.nn.functional. straight_through, so the producer of x trains as if the quantisation were the identity. The codebook itself receives no gradient from that path, which is why the layer also returns the two terms that train it — see VectorQuantizerOutput.