Struct ParallelEmbedding 

pub struct ParallelEmbedding { /* private fields */ }

Vocabulary-partitioned embedding table.

The vocabulary is split evenly across n_workers. Each token index is routed to worker index / (vocab_size / n_workers) and the corresponding row is returned.

Implementations

impl ParallelEmbedding

pub fn new( vocab_size: usize, embed_dim: usize, n_workers: usize, seed: u64, ) -> NeuralResult<Self>

Create a parallel embedding table.

Errors

• vocab_size is not divisible by n_workers.
• n_workers == 0.

pub fn forward(&self, indices: &[usize]) -> NeuralResult<Array2<f64>>

Look up embeddings for a sequence of token indices.

Returns an array of shape [len(indices), embed_dim].

Errors

Returns NeuralError::InvalidArgument if any index >= vocab_size.

pub fn vocab_size(&self) -> usize

Total vocabulary size.

pub fn embed_dim(&self) -> usize

Embedding dimension.