hpt 0.1.3

High Performance Tensor (HPT) - A fast, efficient, and user-friendly tensor computation library for Rust
Documentation
use std::{cell::RefCell, rc::Rc};

use crate::tensor::{DiffTensor, Tensor};
use hpt_allocator::{
    traits::{Allocator, AllocatorOutputRetrive},
    Cpu,
};
use hpt_common::error::base::TensorError;
use hpt_iterator::{iterator_traits::ParStridedIteratorZip, TensorIterator};
use hpt_traits::{
    ops::{creation::TensorCreator, slice::Slice},
    tensor::CommonBounds,
};
use num::Integer;
use rayon::iter::ParallelIterator;
impl<T, const DEVICE: usize, Al> Slice for DiffTensor<T, Cpu, DEVICE, Al>
where
    T: CommonBounds,
    Al: Allocator + 'static + Send + Sync,
    Al::Output: AllocatorOutputRetrive,
{
    fn slice(
        &self,
        index: &[(i64, i64, i64)],
    ) -> Result<DiffTensor<T, Cpu, DEVICE, Al>, TensorError> {
        let res = self.inner.slice(index)?;
        let lhs = self.clone();
        if let None = lhs.grad.borrow().as_ref() {
            lhs.grad.replace(Some(self.inner.zeros_like()?));
        }
        lhs.out_degree.borrow_mut().inc();
        let index = index.to_vec();
        Ok(DiffTensor {
            inner: res,
            grad: Rc::new(RefCell::new(None)),
            out_degree: Rc::new(RefCell::new(0)),
            backward: Rc::new(RefCell::new(move |grad: Tensor<T, Cpu, DEVICE, Al>| {
                let taked = lhs.grad.take();
                if let Some(tmp) = taked {
                    let mut sliced = tmp.slice(&index)?;
                    sliced
                        .par_iter_mut()
                        .zip(grad.inner.par_iter())
                        .for_each(|(a, b)| {
                            *a = a._add(b);
                        });
                    if *lhs.out_degree.borrow() > 1 {
                        *lhs.out_degree.borrow_mut() -= 1;
                    } else {
                        lhs.backward.borrow_mut()(tmp)?;
                    }
                } else {
                    panic!("Gradient is not set for slice");
                }
                Ok(false)
            })),
        })
    }
}