OA LibraryReference

oa::VectorQuantizer

VectorQuantizer class in the OA Ml public surface.

ClC++ public headers

Inheritance

public Module

Public Methods

const oa::Matrix & oa::VectorQuantizer::codebook()

oa::Matrix & oa::VectorQuantizer::codebook()

const VectorQuantizerConfig & oa::VectorQuantizer::config()

void oa::VectorQuantizer::emaUpdate(const oa::Matrix & inZe, const oa::Matrix & inIdx)

oa::Matrix oa::VectorQuantizer::forward(const oa::Matrix & inZe)

oa::Matrix oa::VectorQuantizer::lookup(const oa::Matrix & inIdx)

VqResult oa::VectorQuantizer::quantize(const oa::Matrix & inZe)

void oa::VectorQuantizer::seed(const oa::Matrix & inLatents)

Constructor & Destructor Documentation

oa::VectorQuantizer::VectorQuantizer(
	const VectorQuantizerConfig & inConfig
)

No public source comment is attached to this declaration.

Parameters

inConfig

const VectorQuantizerConfig &

—

void oa::VectorQuantizer::~VectorQuantizer()

No public source comment is attached to this declaration.

Public Method Documentation

const oa::Matrix & oa::VectorQuantizer::codebook()

No public source comment is attached to this declaration.

Returns

const oa::Matrix &

The declared return value.

oa::Matrix & oa::VectorQuantizer::codebook()

No public source comment is attached to this declaration.

Returns

oa::Matrix &

The declared return value.

const VectorQuantizerConfig & oa::VectorQuantizer::config()

No public source comment is attached to this declaration.

Returns

const VectorQuantizerConfig &

The declared return value.

void oa::VectorQuantizer::emaUpdate(
	const oa::Matrix & inZe,
	const oa::Matrix & inIdx
)

EMA codebook update + dead-code reinit. call ONCE per step AFTER the optimizer step, with this step's z_e and the idx returned by quantize(). in-place on the codebook buffer (so it stays the same checkpointed buffers() entry).

Parameters

inZe

const oa::Matrix &

—

inIdx

const oa::Matrix &

—

Returns

void

The declared return value.

oa::Matrix oa::VectorQuantizer::forward(
	const oa::Matrix & inZe
)

Module forward returns just the straight-through quantized tensor.

Parameters

inZe

const oa::Matrix &

—

Returns

oa::Matrix

The declared return value.

oa::Matrix oa::VectorQuantizer::lookup(
	const oa::Matrix & inIdx
)

Token → latent: gather the code vectors for the given indices ([N] Int32, the dtype VqAssign emits, or any generated ids). The inference-time inverse of the nearest-code assignment — feed generated token ids straight back through it to reconstruct z_q [N, D] for the decoder. Pure lookup, no STE.

Parameters

inIdx

const oa::Matrix &

—

Returns

oa::Matrix

The declared return value.

VqResult oa::VectorQuantizer::quantize(
	const oa::Matrix & inZe
)

z_e [N, D] (RMS-normalized latents recommended) → quantized (STE) + code indices + commitment loss. Records entirely on-GPU.

Parameters

inZe

const oa::Matrix &

—

Returns

VqResult

The declared return value.

void oa::VectorQuantizer::seed(
	const oa::Matrix & inLatents
)

Data-dependent init: seed the K codes from the K HIGHEST-NORM rows of inLatents ([>= K, D] encoder outputs) — NOT the first K, which is degenerate for residual VQ (rows a shallow level used → ~zero residual → a deeper codebook seeds all zeros and dies). Highest-norm rows are never zero, so every level gets live, distinct codes. writes the codebook IN-PLACE (copyFrom) so the registered buffers() entry stays valid. Completes inLatents before reading.

Parameters

inLatents

const oa::Matrix &

—

Returns

void

The declared return value.