feat: add 2D and fft model

Author: dancixx

src/attend.rs

@@ -0,0 +1,73 @@
+use candle_core::{DType, Result, Tensor, D};
+
+#[derive(Debug)]
+pub struct Attend {
+    scale: Option<f64>,
+    drop_p: f64,
+    #[allow(dead_code)]
+    heads: usize,
+    is_flash: bool,
+    is_causal: bool,
+}
+
+impl Attend {
+    pub fn new(
+        drop_p: Option<f64>,
+        heads: Option<usize>,
+        scale: Option<f64>,
+        is_flash: Option<bool>,
+        is_causal: Option<bool>,
+    ) -> Self {
+        Self {
+            scale,
+            drop_p: drop_p.unwrap_or(0.0),
+            heads: heads.unwrap_or(8),
+            is_flash: is_flash.unwrap_or(false),
+            is_causal: is_causal.unwrap_or(false),
+        }
+    }
+
+    pub fn forward_t(&self, q: &Tensor, k: &Tensor, v: &Tensor, train: bool) -> Result<Tensor> {
+        let q_dims = q.dims();
+        let scale = self.scale.unwrap_or((q_dims[q_dims.len() - 1] as f64).sqrt());
+
+        if self.is_flash {
+            // Fallback to standard attention for now
+        }
+
+        let q = q.contiguous()?;
+        let k_t = k.transpose(D::Minus2, D::Minus1)?.contiguous()?;
+        let mut sim = q.matmul(&k_t)?;
+        sim = (sim * (1.0 / scale))?;
+
+        if self.is_causal {
+            let sim_dims = sim.dims();
+            let (i, j) = (sim_dims[sim_dims.len() - 2], sim_dims[sim_dims.len() - 1]);
+            // upper triangular mask
+            let mut mask_data = vec![0u8; i * j];
+            for row in 0..i {
+                for col in 0..j {
+                    if col as i64 > row as i64 + (j as i64 - i as i64) {
+                        mask_data[row * j + col] = 1;
+                    }
+                }
+            }
+            let mask = Tensor::from_vec(mask_data, (i, j), sim.device())?
+                .to_dtype(DType::F32)?;
+            let mask_value = f32::NEG_INFINITY;
+            let mask = mask.broadcast_as(sim.dims())?;
+            sim = sim.where_cond(&mask.eq(0.0)?, &Tensor::new(mask_value, sim.device())?)?;
+        }
+
+        let attn = candle_nn::ops::softmax_last_dim(&sim)?;
+        let attn = if train && self.drop_p > 0.0 {
+            candle_nn::ops::dropout(&attn, self.drop_p as f32)?
+        } else {
+            attn
+        };
+
+        let attn = attn.contiguous()?;
+        let v = v.contiguous()?;
+        attn.matmul(&v)
+    }
+}

src/attention.rs

@@ -0,0 +1,218 @@
+use candle_core::{IndexOp, Result, Tensor};
+use candle_nn::{
+    layer_norm, linear_no_bias, LayerNorm, LayerNormConfig, Linear, Module, VarBuilder,
+};
+
+use crate::attend::Attend;
+
+#[derive(Debug)]
+pub struct ToQKV {
+    pub(crate) linear: Linear,
+    pub(crate) heads: usize,
+}
+
+impl ToQKV {
+    pub fn new(vb: VarBuilder, dim: usize, hidden_size: usize, heads: usize) -> Result<Self> {
+        Ok(Self {
+            linear: linear_no_bias(dim, hidden_size * 3, vb)?,
+            heads,
+        })
+    }
+
+    pub fn rearrange(&self, xs: &Tensor) -> Result<(Tensor, Tensor, Tensor)> {
+        let xs_dims = xs.dims();
+        let (h, qkv) = (self.heads, 3);
+        let b = xs_dims[0];
+        let n = xs_dims[1];
+        let total_dim = xs_dims[2];
+        let dim_head = total_dim / (qkv * h);
+        let xs = xs.reshape((b, n, qkv, h, dim_head))?;
+        let xs = xs.permute((2, 0, 3, 1, 4))?;
+
+        let q = xs.i(0)?;
+        let k = xs.i(1)?;
+        let v = xs.i(2)?;
+
+        Ok((q, k, v))
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<(Tensor, Tensor, Tensor)> {
+        let xs = self.linear.forward(xs)?;
+        self.rearrange(&xs)
+    }
+}
+
+#[derive(Debug)]
+pub struct ToValueResidualMix {
+    pub(crate) linear: Linear,
+}
+
+impl ToValueResidualMix {
+    pub fn new(vb: VarBuilder, dim: usize, heads: usize) -> Result<Self> {
+        Ok(Self {
+            linear: linear_no_bias(dim, heads, vb)?,
+        })
+    }
+
+    pub fn rearrange(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = xs.transpose(1, 2)?;
+        xs.unsqueeze(candle_core::D::Minus1)
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = self.linear.forward(xs)?;
+        let xs = self.rearrange(&xs)?;
+        candle_nn::ops::sigmoid(&xs)
+    }
+}
+
+#[derive(Debug)]
+pub struct ToVGates {
+    pub(crate) linear: Linear,
+    #[allow(dead_code)]
+    pub(crate) heads: usize,
+}
+
+impl ToVGates {
+    pub fn new(vb: VarBuilder, dim: usize, heads: usize) -> Result<Self> {
+        Ok(Self {
+            linear: linear_no_bias(dim, heads, vb)?,
+            heads,
+        })
+    }
+
+    pub fn rearrange(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = xs.transpose(1, 2)?;
+        xs.unsqueeze(candle_core::D::Minus1)
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = self.linear.forward(xs)?;
+        let xs = candle_nn::ops::sigmoid(&xs)?;
+        self.rearrange(&xs)
+    }
+}
+
+#[derive(Debug)]
+pub struct ToOut {
+    pub(crate) drop_p: f64,
+    pub(crate) linear: Linear,
+}
+
+impl ToOut {
+    pub fn new(
+        vb: VarBuilder,
+        dim: usize,
+        heads: usize,
+        dim_head: usize,
+        drop_p: Option<f64>,
+    ) -> Result<Self> {
+        Ok(Self {
+            drop_p: drop_p.unwrap_or(0.0),
+            linear: linear_no_bias(dim_head * heads, dim, vb)?,
+        })
+    }
+
+    pub fn rearrange(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs_dims = xs.dims();
+        let (b, h, n, d) = (xs_dims[0], xs_dims[1], xs_dims[2], xs_dims[3]);
+        xs.permute((0, 2, 1, 3))?.reshape((b, n, h * d))
+    }
+
+    pub fn forward_t(&self, xs: &Tensor, train: bool) -> Result<Tensor> {
+        let xs = self.rearrange(xs)?;
+        let xs = self.linear.forward(&xs)?;
+        if train && self.drop_p > 0.0 {
+            candle_nn::ops::dropout(&xs, self.drop_p as f32)
+        } else {
+            Ok(xs)
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct Attention {
+    #[allow(dead_code)]
+    scale: f64,
+    #[allow(dead_code)]
+    drop_p: f64,
+    norm: LayerNorm,
+    pub(crate) to_qkv: ToQKV,
+    pub(crate) to_value_residual_mix: Option<ToValueResidualMix>,
+    pub(crate) to_v_gates: ToVGates,
+    attend: Attend,
+    to_out: ToOut,
+    #[allow(dead_code)]
+    learned_value_residual_mix: bool,
+}
+
+impl Attention {
+    pub fn new(
+        vb: VarBuilder,
+        dim: usize,
+        dim_head: Option<usize>,
+        heads: Option<usize>,
+        drop_p: Option<f64>,
+        is_flash: Option<bool>,
+        learned_value_residual_mix: Option<bool>,
+    ) -> Result<Self> {
+        let dim_head = dim_head.unwrap_or(32);
+        let heads = heads.unwrap_or(4);
+        let scale = (dim_head as f64).sqrt();
+
+        let norm = layer_norm(dim, LayerNormConfig::default(), vb.pp("norm"))?;
+        let to_qkv = ToQKV::new(vb.pp("to_qkv"), dim, dim_head * heads, heads)?;
+        let to_value_residual_mix = if learned_value_residual_mix.unwrap_or(false) {
+            Some(ToValueResidualMix::new(
+                vb.pp("to_value_residual_mix"),
+                dim,
+                heads,
+            )?)
+        } else {
+            None
+        };
+        let to_v_gates = ToVGates::new(vb.pp("to_v_gates"), dim, heads)?;
+        let to_out = ToOut::new(vb.pp("to_out"), dim, heads, dim_head, drop_p)?;
+
+        Ok(Self {
+            scale,
+            drop_p: drop_p.unwrap_or(0.0),
+            norm,
+            to_qkv,
+            to_value_residual_mix,
+            to_v_gates,
+            attend: Attend::new(drop_p, None, None, is_flash, None),
+            to_out,
+            learned_value_residual_mix: learned_value_residual_mix.unwrap_or(false),
+        })
+    }
+
+    pub fn forward_t(
+        &self,
+        xs: &Tensor,
+        value_residual: Option<&Tensor>,
+        train: bool,
+    ) -> Result<(Tensor, Tensor)> {
+        let xs = self.norm.forward(xs)?;
+        let (q, k, mut v) = self.to_qkv.forward(&xs)?;
+        let cache_v = v.clone();
+
+        if let Some(ref to_value_residual_mix) = self.to_value_residual_mix {
+            if let Some(value_residual) = value_residual {
+                let mix = to_value_residual_mix.forward(&xs)?;
+                let diff = value_residual.sub(&v)?;
+                let mix = mix.broadcast_as(diff.dims())?;
+                let weighted = diff.mul(&mix)?;
+                v = v.add(&weighted)?;
+            }
+        }
+
+        let out = self.attend.forward_t(&q, &k, &v, train)?;
+        let gates = self.to_v_gates.forward(&xs)?;
+        let gates = gates.broadcast_as(out.dims())?;
+        let out = out.mul(&gates)?;
+        let out = self.to_out.forward_t(&out, train)?;
+
+        Ok((out, cache_v))
+    }
+}

src/feedforward.rs

@@ -0,0 +1,60 @@
+use candle_core::{IndexOp, Result, Tensor};
+use candle_nn::{layer_norm, linear, LayerNorm, LayerNormConfig, Linear, Module, VarBuilder};
+
+#[derive(Debug, Clone)]
+pub struct GEGLU;
+
+impl GEGLU {
+    pub fn rearrange(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> {
+        let xs_dims = xs.dims();
+        let (b, n, d) = (xs_dims[0], xs_dims[1], xs_dims[2]);
+        let reshaped = xs.reshape((b, n, 2, d / 2))?;
+        let x = reshaped.i((.., .., 0, ..))?;
+        let gate = reshaped.i((.., .., 1, ..))?;
+        Ok((x, gate))
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let (x, gate) = self.rearrange(xs)?;
+        let gate = gate.gelu()?;
+        x.mul(&gate)
+    }
+}
+
+#[derive(Debug)]
+pub struct FeedForward {
+    drop_p: f64,
+    norm: LayerNorm,
+    linear1: Linear,
+    geglu: GEGLU,
+    linear2: Linear,
+}
+
+impl FeedForward {
+    pub fn new(vb: VarBuilder, dim: usize, mult: usize, drop_p: Option<f64>) -> Result<Self> {
+        let hidden_size = ((dim * mult * 2) as f64 / 3.0).trunc() as usize;
+        let norm = layer_norm(dim, LayerNormConfig::default(), vb.pp("norm"))?;
+        let linear1 = linear(dim, hidden_size * 2, vb.pp("linear1"))?;
+        let linear2 = linear(hidden_size, dim, vb.pp("linear2"))?;
+
+        Ok(Self {
+            drop_p: drop_p.unwrap_or(0.0),
+            norm,
+            linear1,
+            geglu: GEGLU,
+            linear2,
+        })
+    }
+
+    pub fn forward_t(&self, xs: &Tensor, train: bool) -> Result<Tensor> {
+        let xs = self.norm.forward(xs)?;
+        let xs = self.linear1.forward(&xs)?;
+        let xs = self.geglu.forward(&xs)?;
+        let xs = if train && self.drop_p > 0.0 {
+            candle_nn::ops::dropout(&xs, self.drop_p as f32)?
+        } else {
+            xs
+        };
+        self.linear2.forward(&xs)
+    }
+}