mojolearn.embedding#

Embedding lookup tables with a reproducible backward pass.

Embedding follows torch.nn.Embedding’s forward and its dense backward. Rows that receive gradient from several positions are summed in ascending position order, so the gradient does not depend on the device or on how the backward is scheduled, and microbatches presented in position order reproduce the unsplit gradient exactly when each one carries the previous gradient forward (backward(ids, dy, grad=...)).

There is no autograd here: the layer holds the table and computes the two calls; it never updates weight itself.

Classes

Embedding(num_embeddings, embedding_dim[, ...])

Embedding table with a reproducible gather and dense gradient, as torch.nn.Embedding.
class mojolearn.embedding.Embedding(num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False, weight=None, plan='auto')[source]#

Bases: NumericModeMixin

Embedding table with a reproducible gather and dense gradient, as torch.nn.Embedding.

The table is the layer’s own float32 copy, taken at construction (and again whenever weight is assigned): to change it, assign layer.weight = new_table; writing into the array you passed in does not reach the layer. On the GPU the table is uploaded once and kept on the device until weight is reassigned.

padding_idx follows torch: a negative value counts from the end, the forward gathers that row like any other, and the backward stores +0.0 in its row. Ids below 0 or at or past num_embeddings raise an error (they are never clamped), as do NaN or infinite values in the table or gradients, and empty ids or dy.

Parameters:

  • num_embeddings (int) – Number of rows V; at least 1.

  • embedding_dim (int) – Row width d; at least 1.

  • padding_idx (int or None, default None) – Row whose gradient is always zero; within [-V, V).

  • max_norm (None) – Only None; anything else raises ValueError.

  • norm_type (float, default 2.0) – Only 2.0 (it has no effect without max_norm).

  • scale_grad_by_freq (bool, default False) – Only False; True raises ValueError.

  • sparse (bool, default False) – Only False (the gradient is the dense (V, d) array); True raises ValueError.

  • weight (array-like of shape (num_embeddings, embedding_dim)) – The table, required (there is no random initialization); copied to float32. Use Embedding.from_pretrained to build a layer from a table alone.

  • plan ({‘auto’, ‘scan’, ‘sort’}, default ‘auto’) – How the backward groups the positions of each row. Both plans give the same bits; ‘auto’ picks by problem size.

Variables:

weight (Array of shape (num_embeddings, embedding_dim), float32) – The table.

Examples

>>> import numpy as np, mojolearn
>>> table = np.random.default_rng(0).random((1000, 64), dtype=np.float32)
>>> emb = mojolearn.Embedding.from_pretrained(table)
>>> ids = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int64)
>>> y = emb.forward(ids)
>>> dW = emb.backward(ids, np.ones((2, 3, 64), dtype=np.float32))
classmethod from_pretrained(embeddings, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False, numeric_mode=None, plan='auto')[source]#

Build a layer from an existing (V, d) table, as torch’s Embedding.from_pretrained.

There is no freeze argument, because nothing here trains the table. The other arguments are the constructor’s.

to_device(a)[source]#

Upload a float32 array once, for use as dy in backward on the device.

forward(ids, device=False)[source]#

Look up rows: weight[ids].

Parameters:

  • ids (array of int) – Row indices, any shape.

  • device (bool, default False) – Return the result as a device tensor, without copying it back to the host. Raises an error when this install has no device path.

Returns:

Array of shape ids.shape + (embedding_dim,), float32

backward(ids, dy, grad=None)[source]#

Dense (num_embeddings, embedding_dim) gradient of forward(ids) for dy.

With grad=None the gradient starts from zero. A grad array is a carried accumulator: the sum continues from its values (it is not modified), so microbatches given in ascending position order reproduce the unsplit gradient bit for bit. A dy from to_device keeps the computation and the result on the device.

Parameters:

  • ids (array of int) – The forward’s ids.

  • dy (array of float32, shape ids.shape + (embedding_dim,)) – Gradient of the forward’s output.

  • grad (array or None, default None) – Gradient to continue from.

Returns:

Array of shape (num_embeddings, embedding_dim), float32

save(path)[source]#

Save the table and settings to path as an .npz file.

A loaded table can look ids up with forward, including on a CPU-only install, where backward is not available.

classmethod load(path)[source]#

Load a layer written by save. The dtype and shape are checked; nothing is cast.