File size: 5,175 Bytes
97df2ae e26f4cc 97df2ae e26f4cc 97df2ae d444358 4dbc907 97df2ae e26f4cc 97df2ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | //! Camadas base: Linear, Embedding, RMSNorm, SwiGLU — port minimalista e
//! autocontido (sem depender do API de módulos do candle-nn para reduzir
//! superfície de drift entre versões).
use candle_core::{DType, Device, Result, Tensor, Var};
use candle_nn::ops::rms_norm;
pub struct Linear {
pub w: Var, // (d_saida, d_entrada)
pub b: Option<Var>,
}
impl Linear {
pub fn novo(d_entrada: usize, d_saida: usize, desvio: f64, dev: &Device) -> Result<Self> {
let w = Tensor::randn(0f32, desvio as f32, (d_saida, d_entrada), dev)?;
let b = Tensor::zeros((d_saida,), DType::F32, dev)?;
Ok(Linear { w: Var::from_tensor(&w)?, b: Some(Var::from_tensor(&b)?) })
}
pub fn sem_vies(d_entrada: usize, d_saida: usize, desvio: f64, dev: &Device) -> Result<Self> {
let w = Tensor::randn(0f32, desvio as f32, (d_saida, d_entrada), dev)?;
Ok(Linear { w: Var::from_tensor(&w)?, b: None })
}
/// v0.9.0 — `sem_vies` com RNG DETERMINÍSTICO próprio (quita o débito
/// r-33 §8.1: a inicialização do JEPA usava o RNG global do candle — não
/// semeado — e impedia bit-identidade entre execuções).
pub fn sem_vies_semeado(
d_entrada: usize,
d_saida: usize,
desvio: f64,
dev: &Device,
rng: &mut crate::rng::RngDeterministico,
) -> Result<Self> {
let n = d_entrada * d_saida;
let dados: Vec<f32> = (0..n).map(|_| rng.normal(0.0, desvio) as f32).collect();
let w = Tensor::from_vec(dados, (d_saida, d_entrada), dev)?;
Ok(Linear { w: Var::from_tensor(&w)?, b: None })
}
/// W = 0 no nascimento — para camadas cujo efeito deve nascer NEUTRO
/// (ex.: W_f do FiLM, Teo R.2: tanh(0·h)=0 ⇒ identidade bit-a-bit).
pub fn zeros_sem_vies(d_entrada: usize, d_saida: usize, dev: &Device) -> Result<Self> {
let w = Tensor::zeros((d_saida, d_entrada), DType::F32, dev)?;
Ok(Linear { w: Var::from_tensor(&w)?, b: None })
}
/// x (..., d_entrada) → (..., d_saida). Achata lotes (candle matmul
/// não aceita 3-D @ 2-D) e devolve na forma original.
pub fn fwd(&self, x: &Tensor) -> Result<Tensor> {
let forma: Vec<usize> = x.dims().to_vec();
let d_saida = self.w.as_tensor().dims()[0];
let x2 = if forma.len() > 2 {
let n = forma.iter().take(forma.len() - 1).product::<usize>();
x.reshape((n, forma[forma.len() - 1]))?
} else {
x.clone()
};
let y = x2.matmul(&self.w.as_tensor().transpose(0, 1)?)?;
let y = match &self.b {
Some(b) => y.broadcast_add(b.as_tensor())?,
None => y,
};
if forma.len() > 2 {
let mut alvo = forma;
*alvo.last_mut().unwrap() = d_saida;
y.reshape(alvo)
} else {
Ok(y)
}
}
pub fn vars(&self) -> Vec<Var> {
match &self.b {
Some(b) => vec![self.w.clone(), b.clone()],
None => vec![self.w.clone()],
}
}
}
pub struct Embedding {
pub tabela: Var, // (vocab, d)
}
impl Embedding {
pub fn novo(vocab: usize, d: usize, dev: &Device) -> Result<Self> {
let t = Tensor::randn(0f32, 0.02, (vocab, d), dev)?;
Ok(Embedding { tabela: Var::from_tensor(&t)? })
}
/// ids u32 (…) → (..., d) — lookup via index_select no plano.
pub fn fwd(&self, ids: &Tensor) -> Result<Tensor> {
let forma: Vec<usize> = ids.dims().to_vec();
let d = self.tabela.as_tensor().dims()[1];
let plano = ids.reshape((ids.elem_count(),))?;
let linhas = self.tabela.as_tensor().index_select(&plano, 0)?;
let mut alvo = forma;
alvo.push(d);
linhas.reshape(alvo)
}
pub fn vars(&self) -> Vec<Var> {
vec![self.tabela.clone()]
}
}
pub struct RmsNorm {
pub ganho: Var,
pub eps: f64,
}
impl RmsNorm {
pub fn novo(d: usize, eps: f64, dev: &Device) -> Result<Self> {
let g = Tensor::ones((d,), DType::F32, dev)?;
Ok(RmsNorm { ganho: Var::from_tensor(&g)?, eps })
}
pub fn fwd(&self, x: &Tensor) -> Result<Tensor> {
rms_norm(x, self.ganho.as_tensor(), self.eps as f32)
}
pub fn vars(&self) -> Vec<Var> {
vec![self.ganho.clone()]
}
}
/// SwiGLU: W2(silu(W1·x) ⊙ W3·x) — d_ff oculto.
pub struct Swiglu {
pub w1: Linear,
pub w3: Linear,
pub w2: Linear,
}
impl Swiglu {
pub fn novo(d: usize, d_ff: usize, dev: &Device) -> Result<Self> {
let desvio = (1.0 / (d as f64).sqrt()) * 0.9;
Ok(Swiglu {
w1: Linear::sem_vies(d, d_ff, desvio, dev)?,
w3: Linear::sem_vies(d, d_ff, desvio, dev)?,
w2: Linear::sem_vies(d_ff, d, desvio / (d_ff as f64).sqrt() * 10.0, dev)?,
})
}
pub fn fwd(&self, x: &Tensor) -> Result<Tensor> {
let a = self.w1.fwd(x)?;
let b = self.w3.fwd(x)?;
let gate = candle_nn::ops::silu(&a)?;
self.w2.fwd(&(&gate * &b)?)
}
pub fn vars(&self) -> Vec<Var> {
[self.w1.vars(), self.w3.vars(), self.w2.vars()].concat()
}
}
|