candela-tensor 0.2.0

A lazy, graph-based tensor engine in 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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

use crate::tensor::errors::OpError;
use crate::tensor::mem_formats::layout::Layout;
use crate::tensor::ops::def_op::OpKind;

#[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip_all))]
pub fn compute_layout<T: Copy>(op: &OpKind<T>, inputs: &[&Layout]) -> Result<Layout, OpError> {
    match op {
        OpKind::ScalarOp(_) | OpKind::FusedScalar(_) => Ok(Layout::new(inputs[0].shape())),
        OpKind::NoOp => Ok(inputs[0].clone()),
        OpKind::View(new_layout)
        | OpKind::Slice(new_layout)
        | OpKind::TransposeAxes(new_layout)
        | OpKind::Broadcast(new_layout) => Ok(new_layout.clone()),
        OpKind::AsContiguous => Ok(Layout::new(inputs[0].shape())),
        OpKind::Transpose => Ok(inputs[0].transpose()),
        OpKind::MatMul(_) => {
            // Assumes that the tensor is ALREADY BROADCASTED!
            let a_shape = inputs[0].shape_as_3d();
            let b_shape = inputs[1].shape_as_3d();

            if a_shape[2] != b_shape[1] {
                return Err(OpError::CannotMatMul(a_shape[2], b_shape[1]));
            };

            if a_shape[0] == 1 && b_shape[0] == 1 {
                return Ok(Layout::new(&[a_shape[1], b_shape[2]]));
            }

            let mut new_shape = inputs[0].shape().to_vec();
            let last = new_shape.len() - 1;
            new_shape[last] = b_shape[2];

            Ok(Layout::new(&new_shape))
        }
        OpKind::MatMulSum(_, _, _) => {
            // Assumes that the tensor is ALREADY BROADCASTED!
            let a_shape = inputs[0].shape_as_3d();
            let b_shape = inputs[1].shape_as_3d();

            if a_shape[2] != b_shape[1] {
                return Err(OpError::CannotMatMul(a_shape[2], b_shape[1]));
            };

            if a_shape[0] == 1 && b_shape[0] == 1 {
                let output_layout = Layout::new(&[a_shape[1], b_shape[2]]);
                if inputs[2].shape() != output_layout.shape() {
                    return Err(OpError::NotSameShape(
                        output_layout.shape().into(),
                        inputs[2].shape().into(),
                    ));
                }

                return Ok(output_layout);
            }

            let mut new_shape = inputs[0].shape().to_vec();
            let last = new_shape.len() - 1;
            new_shape[last] = b_shape[2];

            let output_layout = Layout::new(&new_shape);
            if inputs[2].shape() != output_layout.shape() {
                return Err(OpError::NotSameShape(
                    output_layout.shape().into(),
                    inputs[2].shape().into(),
                ));
            }

            Ok(output_layout)
        }
        OpKind::Add | OpKind::Sub | OpKind::Mul | OpKind::Div => {
            if inputs[0].shape() == inputs[1].shape() {
                Ok(inputs[0].clone())
            } else {
                Err(OpError::NotSameShape(
                    inputs[0].shape().into(),
                    inputs[1].shape().into(),
                ))
            }
        }
        OpKind::Sum | OpKind::Max | OpKind::Mean => Ok(Layout::new(&[1])),
        OpKind::SumAxis(axis, keepdims)
        | OpKind::MaxAxis(axis, keepdims)
        | OpKind::MeanAxis(axis, keepdims) => {
            let axis = if *axis < 0 {
                inputs[0].shape().len() as isize + axis
            } else {
                *axis
            } as usize;

            if axis < inputs[0].shape().len() {
                let mut shape = inputs[0].shape().to_vec();

                if *keepdims {
                    shape[axis] = 1;
                } else {
                    shape.remove(axis);
                }

                // Preserve the rank >= 1 invariant: collapsing the last axis
                // of a 1-D input with keepdims=false would produce [] (0-D).
                if shape.is_empty() {
                    shape.push(1);
                }

                Ok(Layout::new(&shape))
            } else {
                Err(OpError::AxesOutOfBounds)
            }
        }
    }
}

#[cfg(test)]
#[path = "impl_layout_tests.rs"]
mod tests;