redstone-ml 0.0.0

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

use crate::ndarray::NdArrayFlags;
use crate::ops::binary_op_add::BinaryOpAdd;
use crate::ops::binary_op_div::BinaryOpDiv;
use crate::ops::binary_op_mul::BinaryOpMul;
use crate::ops::binary_op_sub::BinaryOpSub;
use crate::ops::binary_ops::{BinaryOpBitAnd, BinaryOpBitOr, BinaryOpRem, BinaryOpShl, BinaryOpShr};
use crate::RawDataType;
use crate::{NdArray, StridedMemory};
use paste::paste;
use std::ops::{AddAssign, BitAndAssign, BitOrAssign, DivAssign, MulAssign, RemAssign, ShlAssign, ShrAssign, SubAssign};


macro_rules! define_binary_iop {
    ( $binary_op_trait:ident, $iop_trait:ident, $operator:tt, $method:ident ) => {
        paste! {
            impl<T: RawDataType + $binary_op_trait> $iop_trait<NdArray<'_, T>> for NdArray<'_, T> {
                fn [<$method _assign>](&mut self, rhs: NdArray<'_, T>) {
                    *self $operator &rhs
                }
            }

            impl<T: RawDataType + $binary_op_trait> $iop_trait<&NdArray<'_, T>> for NdArray<'_, T> {
                fn [<$method _assign>](&mut self, rhs: &NdArray<'_, T>) {
                    if !self.flags.contains(NdArrayFlags::Writeable) {
                        panic!("tensor is readonly.");
                    }
                    
                    if rhs.shape() == self.shape() {
                        unsafe {
                            <T as $binary_op_trait>::$method(self.ptr(), &self.stride(),
                                                             rhs.ptr(), &rhs.stride(),
                                                             self.mut_ptr(), self.shape());
                        }
                    }
                    
                    // right-hand term needs broadcasting
                    else {
                        let rhs = rhs.broadcast_to(&self.shape);
                        
                        unsafe {
                        <T as $binary_op_trait>::$method(self.ptr(), &self.stride(),
                                                         rhs.ptr(), &rhs.stride(),
                                                         self.mut_ptr(), self.shape());
                        }
                    }
                }
            }

            impl<T: RawDataType + $binary_op_trait> $iop_trait<T> for NdArray<'_, T> {
                fn [<$method _assign>](&mut self, rhs: T) {
                    if !self.flags.contains(NdArrayFlags::Writeable) {
                        panic!("tensor is readonly.");
                    }

                    unsafe {
                        <T as $binary_op_trait>::[<$method _scalar>](self.ptr(), self.shape(), self.stride(),
                                                                     rhs, self.mut_ptr());
                    }
                }
            }
        }
    };
}

define_binary_iop!(BinaryOpAdd, AddAssign, +=, add);
define_binary_iop!(BinaryOpSub, SubAssign, -=, sub);
define_binary_iop!(BinaryOpMul, MulAssign, *=, mul);
define_binary_iop!(BinaryOpDiv, DivAssign, /=, div);
define_binary_iop!(BinaryOpRem, RemAssign, %=, rem);
define_binary_iop!(BinaryOpBitAnd, BitAndAssign, &=, bitand);
define_binary_iop!(BinaryOpBitOr, BitOrAssign, |=, bitor);
define_binary_iop!(BinaryOpShl, ShlAssign, <<=, shl);
define_binary_iop!(BinaryOpShr, ShrAssign, >>=, shr);