ella-tensor 0.1.5

Tensor API for the ella datastore.
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

mod binary_arith;
mod boolean;
mod builtin_arith;
mod cmp;
mod combine;
mod constructors;
mod convert;
mod index;
mod masked;
mod reduce;
mod scatter;
mod shape;
mod slice;
mod unary_arith;

use crate::{shape::NdimMax, Shape, Tensor, TensorValue};
use ella_common::ops::{TensorOp, TensorUnaryOp};

fn unary_op<T, O, S, F>(t: &Tensor<T, S>, f: F) -> Tensor<O, S>
where
    T: TensorValue,
    O: TensorValue,
    S: Shape,
    F: Fn(T) -> O,
{
    unsafe { Tensor::from_trusted_len_iter(t.iter().map(f), t.shape().clone()) }
}

fn binary_op<T1, T2, O, S1, S2, F>(
    a: &Tensor<T1, S1>,
    b: &Tensor<T2, S2>,
    op: F,
) -> Tensor<O, <S1 as NdimMax<S2>>::Output>
where
    T1: TensorValue,
    T2: TensorValue,
    O: TensorValue,
    S1: Shape + NdimMax<S2>,
    S2: Shape,
    F: Fn(T1, T2) -> O,
{
    if a.ndim() == b.ndim() && a.shape().slice() == b.shape().slice() {
        let shape = <<S1 as NdimMax<S2>>::Output as Shape>::from_shape(a.shape()).unwrap();
        unsafe {
            Tensor::from_trusted_len_iter(a.iter().zip(b.iter()).map(|(a, b)| op(a, b)), shape)
        }
    } else {
        let (a, b) = a.broadcast_with(b).unwrap();
        let shape = a.shape().clone();
        unsafe {
            Tensor::from_trusted_len_iter(a.iter().zip(b.iter()).map(|(a, b)| op(a, b)), shape)
        }
    }
}