moe-llm-core 1.3.6

Part of the MoE-13 Ternary Intelligence Stack
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

use candle_core::{Result, Tensor};
use candle_nn::{Embedding, Module, VarBuilder};

pub struct EmbeddingLayer {
    embeddings: Embedding,
}

impl EmbeddingLayer {
    pub fn new(vocab_size: usize, hidden_size: usize, vb: VarBuilder) -> Result<Self> {
        let embeddings = candle_nn::embedding(vocab_size, hidden_size, vb.pp("weight"))?;
        Ok(Self { embeddings })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        self.embeddings.forward(x)
    }
}