Embedding

Lookup table and token embedding layer.

class taktiny.nn.Embedding(num_embeddings, embed_features, *, dtype=None, rngs, initializer=<function normal.<locals>.init>, quant=None, axis_names=None, partition_spec=None, metadata=None, precision=None, preferred_element_type=None)[source]

Bases: Module

Stores and looks up embeddings from an N-dimensional table.

The embedding table has shape (*num_embeddings, *embed_features). When num_embeddings is an integer, indices follows conventional embedding semantics and an input of shape (...) produces (..., *embed_features).

For an N-dimensional vocabulary, the final axis of indices contains a coordinate for every vocabulary axis. An input of shape (..., len(num_embeddings)) therefore produces (..., *embed_features).

Parameters:

  • num_embeddings (int | Sequence[int]) – Size of the vocabulary axes. An integer is treated as a one-dimensional vocabulary.

  • embed_features (int | Sequence[int]) – Size of the embedding feature axes. An integer is treated as a one-dimensional feature shape.

  • dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to the table initializer.

  • rngs (Rngs) – Random number generator used to initialize the table.

  • initializer (Initializer) – Function used to initialize the table. Defaults to a normal distribution with standard deviation 0.02.

  • quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for the table.

  • axis_names (tuple[str | None, ...] | None) – Optional logical names for every vocabulary and embedding axis.

  • partition_spec (P | None) – Optional partition specification for the table.

  • metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to the embedding parameter.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Precision reserved for embedding projection operations.

  • preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred result type reserved for embedding projection operations.

Variables:

embedding – The learnable embedding-table parameter.

Examples

Look up vectors from a conventional embedding table:

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> embedding = nn.Embedding(8, 4, rngs=nn.Rngs(0))
>>> embedding(jnp.asarray([1, 3])).shape
(2, 4)

Look up feature matrices using two-dimensional coordinates:

>>> embedding = nn.Embedding((2, 3), (4, 5), rngs=nn.Rngs(1))
>>> coordinates = jnp.asarray([[0, 1], [1, 2]])
>>> embedding(coordinates).shape
(2, 4, 5)