luma-tensor 0.0.1

tensor crate with autograd
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

use crate::{Device, Dim, Float, FloatMeta, Tensor, Storage};

impl<D: Device> Tensor<D, Float> {
    /// Softmax over `dim`. Records `Op::Softmax`.
    pub fn softmax<Dm: Dim>(&self, dim: Dm) -> crate::Result<Self> {
        let dim = dim.to_index(self.shape(), "softmax")?;
        let storage = D::f_softmax(&*self.storage_read()?, self.layout(), dim)?;
        let meta = FloatMeta::on_softmax(self, dim);
        assert_eq!(self.dtype(), storage.dtype());
        Ok(Self::from_storage(storage, self.shape().clone(), meta))
    }

    /// RMSNorm over the last dim: `x / sqrt(mean(x^2)+eps) * weight`.
    /// Records `Op::RmsNorm`.
    pub fn rms_norm(&self, weight: &Self, eps: f64) -> crate::Result<Self> {
        let storage = D::f_rms_norm(&*self.storage_read()?, self.layout(), &*weight.storage_read()?, weight.layout(), eps)?;
        let meta = FloatMeta::on_rms_norm(self, weight, eps);
        assert_eq!(self.dtype(), storage.dtype());
        Ok(Self::from_storage(storage, self.shape().clone(), meta))
    }
}