hpt 0.1.3

High Performance Tensor (HPT) - A fast, efficient, and user-friendly tensor computation library for Rust
Documentation
use crate::tensor_base::_Tensor;
use hpt_allocator::{
    traits::{Allocator, AllocatorOutputRetrive},
    Cpu,
};
use hpt_iterator::{iterator_traits::ParStridedIteratorSimdZip, TensorIterator};
use hpt_traits::ops::creation::TensorCreator;
use hpt_traits::{
    ops::regularization::RegularizationOps,
    tensor::{CommonBounds, TensorInfo},
};
use hpt_types::{
    dtype::TypeCommon,
    into_scalar::Cast,
    traits::SimdSelect,
    type_promote::{Cmp, NormalOutUnary, SimdCmp},
};
use hpt_types::{traits::VecTrait, type_promote::NormalOut};
use rand_distr::Distribution;
use rayon::iter::ParallelIterator;

impl<T, const DEVICE: usize, A> RegularizationOps for _Tensor<T, Cpu, DEVICE, A>
where
    T: CommonBounds + Cmp<Output = bool>,
    T::Vec: SimdCmp,
    <T::Vec as SimdCmp>::Output: SimdSelect<T::Vec>,
    A: Allocator + Send + Sync,
    A::Output: AllocatorOutputRetrive,
{
    type Output = _Tensor<T, Cpu, DEVICE, A>;

    type OutputMeta = T;

    fn dropout(&self, rate: f64) -> Result<Self::Output, hpt_common::error::base::TensorError>
    where
        f64: hpt_types::into_scalar::Cast<Self::OutputMeta>,
        bool: hpt_types::into_scalar::Cast<Self::OutputMeta>,
        Self::OutputMeta: hpt_types::type_promote::NormalOut<bool, Output = Self::OutputMeta>,
    {
        let mut ret = Self::Output::empty(self.shape())?;
        let bernoli = rand_distr::Bernoulli::new(rate)
            .expect("Failed to create Bernoulli distribution for dropout");
        let scale: T = (1.0 / (1.0 - rate)).cast();
        ret.par_iter_mut_simd()
            .zip(self.par_iter_simd())
            .for_each_init(
                || rand::rng(),
                |rng, (ret, val)| {
                    let mask: Self::OutputMeta = bernoli.sample(rng).cast();
                    *ret = val._mul(mask)._mul(scale);
                },
            );
        Ok(ret)
    }

    fn shrinkage(
        &self,
        bias: Self::OutputMeta,
        lambda: Self::OutputMeta,
    ) -> Result<Self::Output, hpt_common::error::base::TensorError> {
        let lambda_vec = <Self::OutputMeta as TypeCommon>::Vec::splat(lambda);
        let bias_vec = <Self::OutputMeta as TypeCommon>::Vec::splat(bias);

        Ok(self
            .par_iter_simd()
            .strided_map_simd(
                |(x, y)| {
                    let shifted = y._sub(bias);
                    let abs_shifted = shifted._abs();
                    let thresholded = abs_shifted._sub(lambda)._max(T::ZERO);
                    *x = shifted._signum()._mul(thresholded);
                },
                |(x, y)| {
                    let shifted = y._sub(bias_vec);
                    let abs_shifted = shifted._abs();
                    let sign_shifted = shifted._signum();

                    let thresholded = abs_shifted._sub(lambda_vec)._max(T::Vec::splat(T::ZERO));

                    x.write_unaligned(sign_shifted._mul(thresholded));
                },
            )
            .collect())
    }
}