Skip to main content

ClikaRT::nn::Embedding

class

Header: ClikaRT/nn/embedding.h
Inherits: ClikaRT::nn::Module

A bound embedding table, the public face of the runtime's embedding lookup, exposed as an nn::Module leaf. ONE class serves both weight postures: a DENSE table (the gather is a pure, dtype-preserving row copy, every dtype a table can carry) and a QUANTIZED table (rows decode at gather; the packed weight is never densely materialized). The posture is read off the bound weight (Tensor::is_quantized()), so consumers hold one Embedding either way.

The weight is the HF nn.Embedding layout: slot weight, [num_embeddings, embedding_dim]. A quantized payload rides a plain Tensor carrying its scheme (a QTensor consumer passes its payload).

The lifecycle (uniform across every weight-bearing nn module): make(...) declares the slot (or binds it, for the from-weight overload); set_weights(...) / load_state_dict(...) bind; the FIRST forward builds the lookup once (thread-safe) and releases the registry slot; the bound primitive is the single resident copy, so named_parameters() no longer lists the packed weight. initialize() is the optional warm-up that pays the build at load time instead of on the first lookup.

Held via std::shared_ptr (an nn::Module leaf); copy/move are pinned by the base.

Embedding-table module: gathers rows of its table by token index, the module form of ops::embedding.

The table is [V, E] (row per vocabulary entry); forward(indices) returns indices.shape ++ [E]. Quantized tables serve through the same module (the gather decodes in place). The module owns its weights: construct with make(...), or declare shapes and bind a checkpoint via load_state_dict. After the first forward (or initialize()) the weight lives ONLY in the backend's packed form (one resident copy); to(dtype) restores, casts, and repacks on the next forward.

auto emb = ClikaRT::nn::Embedding::make(table);
auto h = emb->forward(token_ids); // [B, S] -> [B, S, E]

Static member functions

make(Tensor, optional<DataType>, optional<Device>)

static std::shared_ptr<Embedding> make(
    Tensor weight,
    std::optional<DataType> dtype = std::nullopt,
    std::optional<Device> device = std::nullopt
)

Defaults: dtype = unset (adopt the weight's), device = unset (the weight's own); make(weight) alone adopts everything.

Declared in ClikaRT/nn/embedding.h, line 76

make(int64_t, int64_t, optional<DataType>, optional<Device>)

static std::shared_ptr<Embedding> make(
    std::int64_t num_embeddings,
    std::int64_t embedding_dim,
    std::optional<DataType> dtype = std::nullopt,
    std::optional<Device> device = std::nullopt
)

Defaults: dtype = unset (adopt at bind), device = unset (CPU); the two extents alone declare the table.

Declared in ClikaRT/nn/embedding.h, line 101

Member functions

set_weights()

void set_weights(Tensor weight)

Bind the declared weight slot positionally. A quantized payload rides the tensor (its scheme selects the posture); the declared placement wins; a dense payload must match the declared [num_embeddings, embedding_dim] shape. With no pinned dtype the slot adopts the payload's dtype (never cast at rest). Re-binding drops the built lookup; the next forward rebuilds. Raises ClikaRT::Error on a geometry mismatch.

Declared in ClikaRT/nn/embedding.h, line 116

~Embedding()

~Embedding() override

Declared in ClikaRT/nn/embedding.h, line 120

initialize_impl()

virtual Result<void> initialize_impl() override

Optional warm-up: build the lookup NOW instead of on the first forward (idempotent, thread-safe). The weight slot must hold a real (loaded) tensor; a still-fake slot is a clean error. After the build the registry slot is released; the bound primitive holds the single resident copy.

Declared in ClikaRT/nn/embedding.h, line 127

to_impl(StreamOrDevice)

virtual Result<void> to_impl(StreamOrDevice where) override

Move to a placement (Device / Stream): moves the slots and re-homes the built lookup on the target (the quantized posture re-packs there).

Declared in ClikaRT/nn/embedding.h, line 132

to_impl(DataType)

virtual Result<void> to_impl(DataType dtype) override

Cast the DENSE table to dtype (the built lookup rebuilds at the new dtype on the next forward). On the QUANTIZED posture a dtype move is Unsupported; the weight stays quantized at rest, and dequantize explicitly if a dense copy is wanted.

Declared in ClikaRT/nn/embedding.h, line 137

forward()

Tensor forward(Tensor ids) const

Gather rows: ids [*, S] (Int32 or Int64) -> [*, S, embedding_dim]. Dense: output dtype = the table's (a bit-identical row copy). Quantized: rows decode at gather; output dtype = the pinned dtype, else the scheme's float target. The first call builds (once, thread-safe); later calls reuse the bound lookup. Raises ClikaRT::Error on a still-fake slot or a bad index dtype.

Declared in ClikaRT/nn/embedding.h, line 146