burn-tensor 0.1.0

This library provides multiple tensor implementations hidden behind an easy to use API that supports reverse mode automatic differentiation. 
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

use crate::tensor::{
    backend::ndarray::{BatchMatrix, NdArrayTensor},
    ops::*,
};
use ndarray::{Dim, Dimension};

impl<P, const D: usize> TensorOpsTranspose<P, D> for NdArrayTensor<P, D>
where
    P: Default + Clone + std::fmt::Debug,
    Dim<[usize; D]>: Dimension,
    Dim<[usize; D]>: Dimension,
{
    fn transpose(&self) -> Self {
        if D > 2 {
            return self.clone();
        }

        let batch_matrix = BatchMatrix::from_ndarray(self.array.clone(), self.shape.clone());

        let arrays = batch_matrix
            .arrays
            .iter()
            .map(|matrix| matrix.t())
            .map(|output| output.into_owned().into_shared())
            .collect();

        let mut shape = self.shape.clone();
        let size0 = shape.dims[D - 2];
        let size1 = shape.dims[D - 1];
        shape.dims[D - 2] = size1;
        shape.dims[D - 1] = size0;

        let output = BatchMatrix::new(arrays, shape.clone());

        Self::from_bmatrix(output)
    }
}