Embedding

Embedding#

class braintrace.nn.Embedding#

A lookup table whose gather is routed through the ETP embedding op.

Drop-in replacement for brainstate.nn.Embedding that performs the table lookup with braintrace.embedding(), which is what makes the table eligible for online-learning trace computation. Indices of any rank are accepted; rank 2 and above are folded into one flat axis before the op (the rank-guarded primitive takes only scalar or (batch,) indices) and the output is unfolded back to (*indices.shape, *embedding_size).

Four of the parent’s options are accepted for signature compatibility but not supported, and are rejected by the constructor: max_norm, freeze, scale_grad_by_freq and padding_idx. Each modifies the lookup or its gradient outside the ETP primitive that online learning traces, so there is no way to honour them without diverging from the brainstate semantics. They are part of the signature so that code written against brainstate.nn.Embedding fails with a clear message at the constructor call rather than with a TypeError about an unexpected keyword — or, as before this validation moved, with a deferred failure at the first forward pass, which under jit can be far from the mistake.

Parameters:

  • num_embeddings (int) – Size of the embedding dictionary. Must be non-negative.

  • embedding_size (int or sequence of int) – Size of each embedding vector.

  • embedding_init (Callable or ArrayLike, optional) – Initializer for the lookup table, of shape (num_embeddings, *embedding_size). Default is LecunUniform().

  • padding_idx (int, optional) – Accepted but not supported; anything other than None raises NotImplementedError from the constructor. Zeroing the gradient of one row happens in the parent’s backward rule, which the ETP primitive replaces.

  • max_norm (float, optional) – Accepted but not supported; anything other than None raises NotImplementedError from the constructor. Renormalizing rows inserts a stop_gradient the trace machinery cannot see through.

  • norm_type (float, optional) – The p of the p-norm used by max_norm. Supported in the sense that it is accepted and stored, but inert: it only has an effect together with max_norm, which is rejected. Default is 2.0.

  • scale_grad_by_freq (bool, optional) – Accepted but not supported; True raises NotImplementedError from the constructor. The inverse-frequency scaling lives in the parent’s backward rule, which the ETP primitive replaces. Default is False.

  • freeze (bool, optional) – Accepted but not supported; True raises NotImplementedError from the constructor. Freezing wraps the table in stop_gradient, which removes the very gradient path online learning traces. Default is False.

  • name (str, optional) – Name of the module.

  • param_type (type, optional) – Parameter state type. Default is brainstate.ParamState.

Variables:

weight (brainstate.ParamState) – The learnable table, of shape (num_embeddings, *embedding_size).

Raises:

NotImplementedError – From the constructor, if max_norm, freeze, scale_grad_by_freq or padding_idx is set to a non-default value. The message names every offending option.

See also

braintrace.embedding

The ETP primitive this layer wraps.

brainstate.nn.Embedding

The upstream layer, which supports all options.

Examples

Look up rows of the table:

>>> import jax.numpy as jnp
>>> import braintrace
>>> layer = braintrace.nn.Embedding(10, 4)
>>> layer(jnp.array([0, 3, 3])).shape
(3, 4)

Index arrays of rank 2 or higher are folded and unfolded automatically:

>>> import jax.numpy as jnp
>>> import braintrace
>>> layer = braintrace.nn.Embedding(10, 4)
>>> layer(jnp.array([[0, 1, 2], [3, 4, 5]])).shape
(2, 3, 4)

An unsupported option fails at the constructor, not at the forward pass:

>>> import braintrace
>>> braintrace.nn.Embedding(10, 4, freeze=True)
Traceback (most recent call last):
    ...
NotImplementedError: braintrace.nn.Embedding does not support: freeze=True. ...
update(indices)#

Look up embeddings through the ETP embedding primitive.

Routing the gather through braintrace.embedding() is what makes the table eligible for online-learning trace computation. Indices of rank 2 or higher are folded into one flat axis before the op (the rank-guarded primitive accepts only scalar or (batch,) indices) and the output is unfolded to (*indices.shape, features).

The unsupported-option check is re-run here, not only in __init__. max_norm, freeze, scale_grad_by_freq and padding_idx are plain public attributes that the parent constructor assigns directly, so a caller can still enable one after construction (layer.freeze = True); without this second gate that would silently produce the wrong semantics instead of an error.

Parameters:

indices (ArrayLike) – Integer token indices of any rank.

Returns:

ArrayLike – The gathered embeddings, of shape (*indices.shape, features).

Raises:

NotImplementedError – If max_norm, freeze, scale_grad_by_freq or padding_idx was enabled by assignment after construction. Passing one to the constructor raises there instead.

Embedding.__init__(num_embeddings, embedding_size, embedding_init=LecunUniform(   scale=1.0, mode='fan_in', in_axis=-2, out_axis=-1, distribution='uniform', rng=RandomState(Array((), dtype=key<fry>) overlaying:   [2233333421 2029709265]), unit=Unit("1") ), padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, freeze=False, name=None, param_type=<class 'brainstate.ParamState'>)#

Build the table and reject the options this layer cannot trace.

The parent constructor runs first so that its own argument validation keeps producing the more specific diagnosis where it has one — an out-of-range padding_idx is still the parent’s ValueError, not an “unsupported option” report, because being out of range is a different mistake.

Every argument is forwarded to brainstate.nn.Embedding unchanged and is documented on the class docstring above.

Raises:

NotImplementedError – If max_norm, freeze, scale_grad_by_freq or padding_idx is set to a non-default value.