embedding

→Tensor

embedding(x: Tensor, weight: Tensor, padding_idx: int | None = None, max_norm: float | None = None, norm_type: float = 2.0, scale_grad_by_freq: bool = False, sparse: bool = False)

source edit

Implementing kernel

Look up rows of an embedding table by integer indices.

A learned embedding table maps integer tokens / categorical features into dense vectors:

\mathrm{out}[i_1, \dots, i_k] = W[\, x[i_1, \dots, i_k]\, ]

where W of shape (num_embeddings, embedding_dim) is the lookup table and x holds integer indices in [0, num_embeddings). Equivalent to a one-hot matmul $\mathrm{onehot}(x) W$ but computed with an indexed gather.

Parameters

xTensor

Integer index tensor of arbitrary shape (*), with entries in [0, num_embeddings). A CPU entry outside that range raises IndexError before the lookup and a Metal one is isolated in the gather (see check_embedding_indices); a float tensor raises TypeError.

weightTensor

Embedding table of shape (num_embeddings, embedding_dim).

padding_idxint= None

If given, the embedding vector at weight[padding_idx] is treated as a padding slot: its gradient is forced to zero so the padding embedding stays at its initialised value (typically a zero vector) throughout training. It must lie in [-num_embeddings, num_embeddings) (a negative one counts from the end), else ValueError.

max_normfloat= None

If given, every row of weight that x looks up and whose

L_p

norm exceeds max_norm is rescaled in place, outside autograd, to norm max_norm before the lookup (with

p

= norm_type). weight stays a leaf; a tensor that saved it for an earlier backward sees the write.

norm_typefloat= 2.0

The

p

value of the

L_p

norm used by max_norm. Default 2.0.

scale_grad_by_freqbool= False

If True, the gradient reaching each row of weight from this call is divided by how many times that row's index occurs in x — useful for highly skewed token distributions. The count is per call: two lookups before one backward are each scaled by their own.

sparsebool= False

Request a sparse gradient w.r.t. weight. Gradients are always dense, so True raises NotImplementedError.

Returns

Tensor

Embedded tensor of shape (*, embedding_dim).

Notes

The backward pass for embedding accumulates gradient contributions from repeated indices via scatter-add — multiple tokens of the same type in a batch correctly sum into the same row of $\partial L / \partial W$ .

Examples

>>> import lucid
>>> from lucid.nn.functional import embedding
>>> table = lucid.randn(10, 4)              # 10 tokens, dim 4
>>> ids = lucid.tensor([[1, 2, 4], [4, 3, 2]], dtype=lucid.int64)
>>> out = embedding(ids, table)
>>> out.shape
(2, 3, 4)
With scale_grad_by_freq, row 1 — looked up three times — gets the
mean of its gradients rather than their sum:
>>> w = lucid.ones(5, 2, requires_grad=True)
>>> ids = lucid.tensor([1, 1, 1, 3], dtype=lucid.int64)
>>> embedding(ids, w, scale_grad_by_freq=True).sum().backward()
>>> w.grad[1].tolist(), w.grad[3].tolist()
([1.0, 1.0], [1.0, 1.0])

Used by 2

>>> import lucid >>> from lucid.nn.functional import embedding >>> table = lucid.randn(10, 4) # 10 tokens, dim 4 >>> ids = lucid.tensor([[1, 2, 4], [4, 3, 2]], dtype=lucid.int64) >>> out = embedding(ids, table) >>> out.shape (2, 3, 4) With scale_grad_by_freq, row 1 — looked up three times — gets the mean of its gradients rather than their sum: >>> w = lucid.ones(5, 2, requires_grad=True) >>> ids = lucid.tensor([1, 1, 1, 3], dtype=lucid.int64) >>> embedding(ids, w, scale_grad_by_freq=True).sum().backward() >>> w.grad[1].tolist(), w.grad[3].tolist() ([1.0, 1.0], [1.0, 1.0])