mojolearn.Embedding

class mojolearn.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')

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))

__init__(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')

Methods

backward(ids, dy[, grad])	Dense (num_embeddings, embedding_dim) gradient of forward(ids) for dy.
forward(ids[, device])	Look up rows: weight[ids].
from_pretrained(embeddings[, padding_idx, ...])	Build a layer from an existing (V, d) table, as torch's Embedding.from_pretrained.
get_params([deep])
load(path)	Load a layer written by save.
numeric_mode_used()	The numeric mode this estimator runs in.
save(path)	Save the table and settings to path as an .npz file.
set_params(**params)	Set constructor parameters.
to_device(a)	Upload a float32 array once, for use as dy in backward on the device.
vendor_used()	The GPU backend this estimator runs on.

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')

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)

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

forward(ids, device=False)

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)

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)

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)

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