mixle.models.embedding module

CategoricalEmbedding – a learned vector per category, usable (and tie-able) in several models at once.

An embedding turns a categorical value – a word/token, a country, a product id – into a learned dense vector. This is a declarative handle for one such embedding table of shape (num_categories, dim): it builds a single nn.Embedding lazily and returns that same module to every model that references it, so passing the same instance to several models ties their vectors and trains them jointly (the neural analogue of the PPL’s name= tying for scalar latents). A word embedding shared across the language-model experts of a mixture is the headline case, but the primitive embeds any categorical field.

Pass a CategoricalEmbedding as embedding= to mixle.models.StreamingTransformerLeaf (.from_config), mixle.models.language_model.LM, mixle.models.transformer.build_causal_lm(), or the PPL Transformer(embedding=...) token. In the PPL it is exposed as mixle.ppl.Embedding.

class CategoricalEmbedding(num_categories, dim, *, name=None)

Bases: object

A lazily-built learned embedding of shape (num_categories, dim); every consumer gets the same module.

Parameters:

• num_categories (int)

• dim (int)

• name (str | None)

module()

The underlying nn.Embedding – built on first call, the identical instance thereafter.

Return type:

Any

resolve_embedding(embedding, num_categories, dim)

Normalize embedding (CategoricalEmbedding | nn.Embedding | None) to an nn.Embedding or None.

Validates that the resolved embedding matches (num_categories, dim) so a shape mismatch fails early with a clear message rather than deep inside a forward pass.

Parameters:

• embedding (Any)

• num_categories (int)

• dim (int)

Return type:

Any