chela 0.0.2

High-performance Machine Learning, Auto-Differentiation and Tensor Algebra crate for Rust
Documentation 
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

use crate::autograd::util::reduce_gradient;
use crate::gradient_function::{GradientFuncTrait, GradientFunction};
use crate::{FloatDataType, NdArray, StridedMemory, Tensor};
use std::cell::RefCell;
use std::rc::Rc;


pub(crate) struct MulBackwards<'a, T: FloatDataType> {
    pub(super) next_functions: [GradientFunction<T>; 2],

    pub(super) lhs_grad: NdArray<'a, T>,
    pub(super) rhs_grad: NdArray<'a, T>,

    pub(super) lhs_shape: Vec<usize>,
    pub(super) rhs_shape: Vec<usize>
}

pub(crate) struct MulScalarBackwards<T: FloatDataType> {
    next_function: GradientFunction<T>,

    shape: Vec<usize>,
    scalar: T,
}


impl<T: FloatDataType> GradientFuncTrait<T> for MulBackwards<'_, T> {
    fn backward(&mut self, grad: &NdArray<T>) {
        let lhs_grad = &self.lhs_grad * grad;
        let rhs_grad = &self.rhs_grad * grad;

        let lhs_grad = reduce_gradient(&lhs_grad, &self.lhs_shape);
        let rhs_grad = reduce_gradient(&rhs_grad, &self.rhs_shape);

        self.next_functions[0].borrow_mut().backward(&lhs_grad);
        self.next_functions[1].borrow_mut().backward(&rhs_grad);
    }
}

impl<T: FloatDataType> GradientFuncTrait<T> for MulScalarBackwards<T> {
    fn backward(&mut self, grad: &NdArray<T>) {
        let grad = grad * self.scalar;
        let grad = reduce_gradient(&grad, &self.shape);

        self.next_function.borrow_mut().backward(&grad);
    }
}


impl<T: FloatDataType> MulBackwards<'static, T> {
    pub(crate) fn new(lhs: &Tensor<T>, rhs: &Tensor<T>) -> GradientFunction<T> {
        let next_functions = [lhs.grad_fn(), rhs.grad_fn()];

        let grad_fn = Self {
            next_functions,
            lhs_grad: rhs.detach(),
            rhs_grad: lhs.detach(),

            lhs_shape: lhs.shape().to_vec(),
            rhs_shape: rhs.shape().to_vec(),
        };

        Rc::new(RefCell::new(grad_fn))
    }
}

impl<T: FloatDataType> MulScalarBackwards<T> {
    pub(crate) fn new(lhs: &Tensor<T>, rhs: T) -> GradientFunction<T> {
        let grad_fn = Self {
            next_function: lhs.grad_fn(),
            shape: lhs.shape().to_vec(),
            scalar: rhs
        };

        Rc::new(RefCell::new(grad_fn))
    }
}