embedding

Contents

embedding#

class braintrace.embedding(indices, weight, *, weight_fn=None)[source]#

ETP-aware embedding lookup (trainable gather).

Computes y = weight_fn(weight)[indices], routed through an ETP primitive so the table participates in eligibility-trace computation. Auto-dispatches on the index rank: etp_emb_p for a (batch,) index vector, etp_emb_v_p for a scalar index.

Parameters:

  • indices (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Integer token indices, scalar () or rank-1 (batch,). The indices are never differentiated.

  • weight (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – The embedding table, of shape (num_embeddings, features). May be a brainunit.Quantity; the unit is split off, the lookup is computed on the mantissa, and the unit is reattached to the result.

  • weight_fn (Callable[[Array | ndarray | bool | number | bool | int | float | complex | Quantity], Array | ndarray | bool | number | bool | int | float | complex | Quantity] | None) – Element-wise transform applied to the table inside the primitive before the lookup (e.g. a mask or normalization). Its Jacobian is composed automatically in the weight-gradient rule.

Returns:

The gathered rows, of shape (features,) for a scalar index or (batch, features) for a (batch,) index vector.

Return type:

Array | ndarray | bool | number | bool | int | float | complex | Quantity

Raises:

  • TypeError – If indices is not of integer dtype.

  • NotImplementedError – If indices.ndim >= 2. Traces are defined per-step, per-batch-element; embed one step at a time or flatten the leading axes outside the op.

  • ValueError – If weight is not a rank-2 (num_embeddings, features) matrix.

See also

matmul

ETP-aware dense matrix multiplication.

grouped_matmul

ETP-aware block-diagonal (grouped) matrix multiplication.

Notes

The D-RTRL eligibility trace for the table has shape (batch, num_embeddings, features, n_state) — it scales linearly with the vocabulary size. For large vocabularies prefer braintrace.pp_prop() (ES-D-RTRL), whose trace is output-shaped ((batch, features, n_state)) and independent of the vocabulary size.

Examples

>>> import jax.numpy as jnp
>>> import braintrace
>>> table = jnp.arange(12, dtype=jnp.float32).reshape(4, 3)
>>> tokens = jnp.array([0, 2], dtype=jnp.int32)
>>> braintrace.embedding(tokens, table)
Array([[0., 1., 2.],
       [6., 7., 8.]], dtype=float32)
>>>
>>> import brainunit as u
>>> y = braintrace.embedding(tokens, table * u.mV)
>>> y.unit
mvolt
Contents