auto_diff/op/
local.rs

1#![allow(clippy::redundant_closure_call)]
2use tensor_rs::tensor::Tensor;
3use super::{OpTrait, OpHandle};
4use super::macros::new_binary_op;
5
6#[cfg(feature = "use-serde")]
7use serde::{Serialize, Deserialize};
8#[cfg(feature = "use-serde")]
9use std::any::Any;
10
11
12
13new_binary_op!(Add, "Add",
14               (|a:&[Tensor], b:&[Tensor]|
15                b[0].swap(&a[0].add(&a[1]))
16               ),
17               (|input: &[Tensor], output_grad: &[Tensor], input_grad: &[Tensor]| {
18                   let x = input[0].ones_like().mul(&output_grad[0]);
19                   let y = input[1].ones_like().mul(&output_grad[0]);
20                   input_grad[0].swap(&x);
21                   input_grad[1].swap(&y);
22               })
23);
24new_binary_op!(Sub, "Sub",
25               (|a:&[Tensor], b:&[Tensor]|
26                b[0].swap(&a[0].sub(&a[1]))),
27               (|input: &[Tensor], output_grad: &[Tensor], input_grad: &[Tensor]| {
28                   let x = input[0].ones_like().mul(&output_grad[0]);
29                   let y = input[1].ones_like().neg().mul(&output_grad[0]);
30                   input_grad[0].swap(&x);
31                   input_grad[1].swap(&y);
32               })
33);
34new_binary_op!(Mul, "Mul",
35               (|a:&[Tensor], b:&[Tensor]|
36                b[0].swap(&a[0].mul(&a[1]))),
37               (|input: &[Tensor], output_grad: &[Tensor], input_grad: &[Tensor]| {
38                   let x = input[1].mul(&output_grad[0]);
39                   let y = input[0].mul(&output_grad[0]);
40                   input_grad[0].swap(&x);
41                   input_grad[1].swap(&y);
42               })
43);
44new_binary_op!(Div, "Div",
45               (|a:&[Tensor], b:&[Tensor]|
46                b[0].swap(&a[0].div(&a[1]))),
47               (|input: &[Tensor], output_grad: &[Tensor], input_grad: &[Tensor]| {
48                   let x = input[1].reciprocal().mul(&output_grad[0]);
49                   let y = input[0].neg().div(&input[1]).div(&input[1]).mul(&output_grad[0]);
50                   input_grad[0].swap(&x);
51                   input_grad[1].swap(&y);
52               })
53);
54
55new_binary_op!(Matmul, "Matmul",
56               (|a:&[Tensor], b:&[Tensor]|
57                b[0].swap(&a[0].matmul(&a[1]))),
58               (|input: &[Tensor], output_grad: &[Tensor], input_grad: &[Tensor]| {
59                   input_grad[0].swap(&input[1].outer(&output_grad[0], Some(true)));
60                   input_grad[1].swap(&input[0].outer(&output_grad[0], Some(true)));
61               })
62);
63
64new_binary_op!(Outer, "Outer",
65               (|a:&[Tensor], b:&[Tensor]|
66                b[0].swap(&a[0].outer(&a[1], None))),
67               (|input: &[Tensor], output_grad: &[Tensor], input_grad: &[Tensor]| {
68                   unimplemented!();
69               })
70);
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75    use crate::op::_gradient_checker;
76
77    #[test]
78    fn matmul() {
79        let mut op = Mul::new();
80
81        for i in 0..10 {
82            let zero = Tensor::from_vec_f64(&vec![(i - 5) as f64], &vec![1]);
83            let zero2 = zero.clone();
84            let good_grad = _gradient_checker(&mut op, &[zero, zero2], None, None, None);
85            assert_eq!(good_grad, true);                        
86        }
87    }
88}
auto_diff/op/local.rs

auto_diff/op/
local.rs
