acme-tensor 0.3.1

A comprehensive tensor library for Rust with support for automatic-differentiation.
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
116
117
118
119

/*
    Appellation: kinds <mod>
    Contrib: FL03 <jo3mccain@icloud.com>
*/
use crate::ops::kinds::reshape::*;
use crate::shape::{Axis, Shape};
use crate::tensor::TensorBase;
use acme::prelude::{BinaryOp, UnaryOp};
use num::Complex;

pub type BoxTensor<T = f64> = Box<TensorBase<T>>;

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
#[non_exhaustive]
pub enum TensorExpr<A, B = A> {
    Binary(BoxTensor<A>, BoxTensor<B>, BinaryOp),
    BinaryScalar(BoxTensor<A>, B, BinaryOp),
    Unary(BoxTensor<A>, UnaryOp),
    Matmul(BoxTensor<A>, BoxTensor<B>),
    Sigmoid(BoxTensor<A>),
    Shape(ReshapeExpr<A>),
}

impl<A, B> TensorExpr<A, B> {
    pub fn binary(lhs: TensorBase<A>, rhs: TensorBase<B>, op: BinaryOp) -> Self {
        Self::Binary(lhs.boxed(), rhs.boxed(), op)
    }

    pub fn binary_scalar(lhs: TensorBase<A>, rhs: B, op: BinaryOp) -> Self {
        Self::BinaryScalar(Box::new(lhs), rhs, op)
    }

    pub fn binary_scalar_c(
        lhs: TensorBase<A>,
        rhs: Complex<A>,
        op: BinaryOp,
    ) -> TensorExpr<A, Complex<A>> {
        TensorExpr::BinaryScalar(Box::new(lhs), rhs, op)
    }

    pub fn broadcast(tensor: TensorBase<A>, shape: Shape) -> Self {
        Self::shape(ReshapeExpr::broadcast(tensor, shape))
    }

    pub fn matmul(lhs: TensorBase<A>, rhs: TensorBase<B>) -> Self {
        Self::Matmul(Box::new(lhs), Box::new(rhs))
    }

    pub fn reshape(tensor: TensorBase<A>, shape: Shape) -> Self {
        Self::shape(ReshapeExpr::reshape(tensor, shape))
    }

    pub fn shape(expr: ReshapeExpr<A>) -> Self {
        Self::Shape(expr)
    }

    pub fn sigmoid(tensor: TensorBase<A>) -> Self {
        Self::Sigmoid(Box::new(tensor))
    }

    pub fn swap_axes(tensor: TensorBase<A>, swap: Axis, with: Axis) -> Self {
        Self::shape(ReshapeExpr::swap_axes(tensor, swap, with))
    }

    pub fn transpose(tensor: TensorBase<A>) -> Self {
        Self::Shape(ReshapeExpr::transpose(tensor))
    }

    pub fn unary(tensor: TensorBase<A>, op: UnaryOp) -> Self {
        Self::Unary(Box::new(tensor), op)
    }

    pub fn lhs(self) -> Option<TensorBase<A>> {
        match self {
            Self::Binary(lhs, _, _) => Some(*lhs),
            Self::BinaryScalar(lhs, _, _) => Some(*lhs),
            Self::Unary(lhs, _) => Some(*lhs),
            Self::Matmul(lhs, _) => Some(*lhs),
            _ => None,
        }
    }

    pub fn rhs(self) -> Option<TensorBase<B>> {
        match self {
            Self::Binary(_, rhs, _) => Some(*rhs),
            Self::BinaryScalar(_, scalar, _) => Some(TensorBase::from_scalar(scalar)),
            Self::Matmul(_, rhs) => Some(*rhs),
            _ => None,
        }
    }
    pub fn view(&self) -> TensorExpr<&A, &B> {
        match self {
            TensorExpr::Binary(lhs, rhs, op) => TensorExpr::binary(lhs.view(), rhs.view(), *op),
            TensorExpr::BinaryScalar(lhs, rhs, op) => {
                TensorExpr::binary_scalar(lhs.view(), rhs, *op)
            }
            TensorExpr::Unary(tensor, op) => TensorExpr::unary(tensor.view(), *op),
            TensorExpr::Matmul(lhs, rhs) => TensorExpr::matmul(lhs.view(), rhs.view()),
            TensorExpr::Sigmoid(tensor) => TensorExpr::sigmoid(tensor.view()),
            TensorExpr::Shape(inner) => TensorExpr::Shape(inner.view()),
        }
    }
    pub fn view_mut(&mut self) -> TensorExpr<&mut A, &mut B> {
        match self {
            TensorExpr::Binary(lhs, rhs, op) => {
                TensorExpr::binary(lhs.view_mut(), rhs.view_mut(), *op)
            }

            TensorExpr::BinaryScalar(lhs, rhs, op) => {
                TensorExpr::binary_scalar(lhs.view_mut(), rhs, *op)
            }
            TensorExpr::Unary(tensor, op) => TensorExpr::unary(tensor.view_mut(), *op),
            TensorExpr::Matmul(lhs, rhs) => TensorExpr::matmul(lhs.view_mut(), rhs.view_mut()),
            TensorExpr::Sigmoid(tensor) => TensorExpr::sigmoid(tensor.view_mut()),
            TensorExpr::Shape(inner) => TensorExpr::Shape(inner.view_mut()),
        }
    }
}