Module Tolk_nn.Embedding

Token embedding: a lookup table from integer indices to dense vectors.

type t = {
  1. weight : Tolk_frontend.Tensor.t;
    (*

    Shape (vocab_size, embed_size).

    *)
}

An embedding layer.

val create : int -> int -> t

create vocab_size embed_size is an embedding table of vocab_size rows of embed_size features, initialised with Rand.glorot_uniform. Load trained values with State.load_state_dict.

apply e idx looks up the row of the table for every element of the integer tensor idx, appending an axis of size embed_size to its shape.

  • raises Invalid_argument

    if idx is not an integer tensor.