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

use crate::StridedMemory;
use crate::broadcast::broadcast_shapes;
use crate::{IntegerDataType, NdArray, RawDataType};
use std::ops::{Add, BitAnd, BitOr, Div, Mul, Rem, Shl, Shr, Sub};

use paste::paste;

macro_rules! define_binary_ops {
    ($object:ident, $($trait_: ident, $operator: tt, $method: ident;)* ) => {
        $(
            fn $method<'a, 'b>(lhs: impl AsRef<$object<'a, T>>,
                               rhs: impl AsRef<$object<'b, T>>) -> $object<'static, T>
            where
                T: $trait_<Output=T>,
            {
                let lhs = lhs.as_ref();
                let rhs = rhs.as_ref();

                let shape = broadcast_shapes(lhs.shape(), rhs.shape());
                let lhs = lhs.broadcast_to(&shape);
                let rhs = rhs.broadcast_to(&shape);

                let data = lhs.flatiter().zip(rhs.flatiter()).map(|(lhs, rhs)| lhs $operator rhs).collect();
                unsafe { $object::from_contiguous_owned_buffer(shape, data) }
            }

            paste! { fn [<$method _scalar>] <'a, 'b>(lhs: impl AsRef<$object<'a, T>>,
                                                     rhs: T) -> $object<'static, T>
                where
                    T: $trait_<Output=T>,
                {
                    let lhs = lhs.as_ref();

                    let data = lhs.flatiter().map(|lhs| lhs $operator rhs).collect();
                    unsafe { $object::from_contiguous_owned_buffer(lhs.shape().to_vec(), data) }
                }
            }
        )*
    }
}

pub(crate) trait BinaryOps<T: RawDataType> {
    define_binary_ops!(
        NdArray,
        Add, +, add;
        Sub, -, sub;
        Mul, *, mul;
        Div, /, div;
        Rem, %, rem;
        BitAnd, &, bitand;
        BitOr, |, bitor;
        Shl, <<, shl;
        Shr, >>, shr;
    );
}

impl<T: IntegerDataType> BinaryOps<T> for T {}
impl BinaryOps<bool> for bool {}
impl BinaryOps<f32> for f32 {}
impl BinaryOps<f64> for f64 {}