Struct Embedding 

pub struct Embedding<B>
where
    B: Backend,
{ /* private fields */ }

A learnable lookup table mapping integer indices to dense vectors.

Examples

let emb = Embedding::<CpuBackend>::new(1000, 128, DType::F32, &dev)?;
// Input: token indices [batch=2, seq_len=5]
let tokens = CpuTensor::from_f64_slice(&indices, (2, 5), DType::I64, &dev)?;
let vectors = emb.forward(&tokens)?; // [2, 5, 128]

Implementations

impl<B> Embedding<B>

where B: Backend,

pub fn new( num_embeddings: usize, embedding_dim: usize, dtype: DType, device: &<B as Backend>::Device, ) -> Result<Embedding<B>, Error>

Create a new Embedding layer with normally-distributed random weights.

pub fn from_tensor(weight: Tensor<B>) -> Result<Embedding<B>, Error>

Create from an existing weight matrix.

pub fn num_embeddings(&self) -> usize

pub fn embedding_dim(&self) -> usize

pub fn weight(&self) -> &Tensor<B>

Trait Implementations

impl<B> Module<B> for Embedding<B>

where B: Backend,

fn forward(&self, indices: &Tensor<B>) -> Result<Tensor<B>, Error>

Look up embeddings for the given indices.

Input: integer tensor of any shape […] Output: tensor of shape […, embedding_dim]

Uses index_select to gather rows from the weight matrix directly on-device (no host round-trip). For autograd, we record the operation so gradients can flow back through the embedding table.

fn parameters(&self) -> Vec<Tensor<B>>

Return all trainable parameters of this module. The optimizer uses these to update weights during training.

fn named_parameters(&self) -> Vec<(String, Tensor<B>)>

Return all trainable parameters with human-readable names.

fn set_training(&self, _training: bool)

Set training or evaluation mode.

fn is_training(&self) -> bool

Whether the module is in training mode (default: true).

fn train(&self)

Convenience: set training mode.

fn eval(&self)

Convenience: set evaluation mode.

fn num_parameters(&self) -> usize

Total number of scalar parameters in this module.

fn trainable_params_count(&self) -> usize

Number of trainable (variable) parameters.

fn frozen_parameters(&self) -> Vec<Tensor<B>>

Freeze all parameters: returns new parameter tensors with is_variable = false, preventing gradient accumulation.

fn state_dict(&self) -> Vec<(String, Tensor<B>)>

Returns a state_dict-style map of parameter name → tensor.