Struct Embedding 

pub struct Embedding {
    pub vocab_size: u64,
    pub embed_size: u64,
    pub weight: Tensor,
    pub arange: Tensor,
}

Embedding layer

Fields

vocab_size: u64

Vocabulary size

embed_size: u64

Embedding size

weight: Tensor

Weight

arange: Tensor

Arange

Implementations

impl Embedding

pub fn new( vocab_size: u64, embed_size: u64, dtype: DType, ) -> Result<Embedding, ZyxError>

new embedding layer

pub fn from_params(weight: Tensor) -> Result<Embedding, ZyxError>

Initialize embedding using only weight

pub fn forward(&self, x: impl Into<Tensor>) -> Result<Tensor, ZyxError>

Forward embedding layer

Trait Implementations

impl Debug for Embedding

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<'a> IntoIterator for &'a Embedding

type Item = &'a Tensor

The type of the elements being iterated over.

type IntoIter = IntoIter<&'a Tensor>

Which kind of iterator are we turning this into?

fn into_iter(self) -> Self::IntoIter

Creates an iterator from a value.

impl<'a> IntoIterator for &'a mut Embedding

type Item = &'a mut Tensor

The type of the elements being iterated over.

type IntoIter = IntoIter<&'a mut Tensor>

Which kind of iterator are we turning this into?

fn into_iter(self) -> Self::IntoIter

Creates an iterator from a value.

impl Module for Embedding

fn iter<'a>(&'a self) -> impl Iterator<Item = &'a Tensor>

Iterate over all tensors immutably

fn iter_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut Tensor>

Iterate over all tensors mutably

fn iter_tensors<'a>(&'a self) -> impl Iterator<Item = (String, &'a Tensor)>

Iterate over tensors without consuming the module

fn iter_tensors_mut<'a>( &'a mut self, ) -> impl Iterator<Item = (String, &'a mut Tensor)>

From tensors

fn realize(&self) -> Result<(), ZyxError>

Realize all tensors in the module

fn set_params(&mut self, params: &mut HashMap<String, Tensor>)

Set parameters, removes them from params, skips parameters that are not found in params.

fn save(&self, path: impl AsRef<Path>) -> Result<(), ZyxError>

Save tensors or modules to a file determined by file extension. Currently only safetensors is supported format.