burn-ndarray 0.18.0

Ndarray backend for the Burn framework
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

use crate::{element::FloatNdArrayElement, tensor::NdArrayTensor};
use burn_tensor::{Element, TensorMetadata};
use ndarray::{Array4, Array5};

use super::NdArrayOps;

pub(crate) fn apply_padding_4d<E: Element>(
    x: NdArrayTensor<E>,
    padding: [usize; 2],
    elem: E,
) -> NdArrayTensor<E> {
    let [batch_size, input_channels, height, width] = x.shape().dims();
    let [padding_height, padding_width] = padding;
    let padded_height = height + 2 * padding_height;
    let padded_width = width + 2 * padding_width;

    let x_new = Array4::from_elem(
        (batch_size, input_channels, padded_height, padded_width),
        elem,
    );
    let mut x_new = NdArrayTensor::new(x_new.into_shared().into_dyn());

    x_new = NdArrayOps::slice_assign(
        x_new,
        &[
            0..batch_size,
            0..input_channels,
            padding_height..height + padding_height,
            padding_width..width + padding_width,
        ],
        x,
    );

    x_new
}

pub(crate) fn apply_padding_5d<E: FloatNdArrayElement>(
    x: NdArrayTensor<E>,
    padding: [usize; 3],
    elem: E,
) -> NdArrayTensor<E> {
    let [batch_size, input_channels, depth, height, width] = x.shape().dims();
    let [padding_depth, padding_height, padding_width] = padding;
    let padded_depth = depth + 2 * padding_depth;
    let padded_height = height + 2 * padding_height;
    let padded_width = width + 2 * padding_width;

    let x_new = Array5::from_elem(
        (
            batch_size,
            input_channels,
            padded_depth,
            padded_height,
            padded_width,
        ),
        elem,
    );
    let mut x_new = NdArrayTensor::new(x_new.into_shared().into_dyn());

    x_new = NdArrayOps::slice_assign(
        x_new,
        &[
            0..batch_size,
            0..input_channels,
            padding_depth..depth + padding_depth,
            padding_height..height + padding_height,
            padding_width..width + padding_width,
        ],
        x,
    );

    x_new
}