native_neural_network_std 0.2.1

Ergonomic std wrapper for the `native_neural_network` crate (no_std) — std-friendly re-exports and utilities.
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

#[derive(Clone, Copy, Debug)]
pub struct Conv3DShapeParams {
    pub kernel: (usize, usize, usize),
    pub pad: (usize, usize, usize),
    pub stride: (usize, usize, usize),
}

pub fn conv3d_output_shape(
    d: usize,
    h: usize,
    w: usize,
    params: Conv3DShapeParams,
) -> Option<(usize, usize, usize)> {
    let (kd, kh, kw) = params.kernel;
    let args = native_neural_network::conv3d::Conv3dOutputShapeArgs {
        d,
        h,
        w,
        kd,
        kh,
        kw,
        pad: params.pad,
        stride: params.stride,
    };
    native_neural_network::conv3d::conv3d_output_shape(args)
}

pub fn conv3d_forward(
    input: &mut crate::std::tensor_std::TensorStd,
    kernel: &mut crate::std::tensor_std::TensorStd,
    bias: Option<&[f32]>,
    output: &mut crate::std::tensor_std::TensorStd,
    pad: (usize, usize, usize),
    stride: (usize, usize, usize),
) -> Result<(), &'static str> {
    if input.shape.len() != 5 || kernel.shape.len() != 5 || output.shape.len() != 5 {
        return Err("tensor rank must be 5");
    }
    let mut in_shape = [0usize; 5];
    for (i, &v) in input.shape.iter().enumerate().take(5) {
        in_shape[i] = v;
    }
    let mut k_shape = [0usize; 5];
    for (i, &v) in kernel.shape.iter().enumerate().take(5) {
        k_shape[i] = v;
    }
    let mut out_shape = [0usize; 5];
    for (i, &v) in output.shape.iter().enumerate().take(5) {
        out_shape[i] = v;
    }

    let in_view = native_neural_network::tensor::TensorView {
        data: input.as_mut_slice(),
        shape: in_shape,
    };
    let k_view = native_neural_network::tensor::TensorView {
        data: kernel.as_mut_slice(),
        shape: k_shape,
    };
    let mut out_view = native_neural_network::tensor::TensorView {
        data: output.as_mut_slice(),
        shape: out_shape,
    };

    native_neural_network::conv3d::conv3d_forward(
        &in_view,
        &k_view,
        bias,
        &mut out_view,
        pad,
        stride,
    );
    Ok(())
}

pub fn conv3d_is_compatible(
    input: &crate::std::tensor_std::TensorStd,
    kernel: &crate::std::tensor_std::TensorStd,
    output: &crate::std::tensor_std::TensorStd,
) -> bool {
    if input.shape.len() != 5 || kernel.shape.len() != 5 || output.shape.len() != 5 {
        return false;
    }

    if input.shape[1] != kernel.shape[1] {
        return false;
    }
    if output.shape[0] != input.shape[0] {
        return false;
    }
    if output.shape[1] != kernel.shape[0] {
        return false;
    }

    let prod =
        |s: &[usize]| -> usize { s.iter().copied().fold(1usize, |a, b| a.saturating_mul(b)) };
    if prod(&input.shape) != input.data.len() {
        return false;
    }
    if prod(&kernel.shape) != kernel.data.len() {
        return false;
    }
    if prod(&output.shape) != output.data.len() {
        return false;
    }
    true
}

pub fn conv3d_layout_compatible(
    input: &crate::std::tensor_std::TensorStd,
    kernel: &crate::std::tensor_std::TensorStd,
    output: &crate::std::tensor_std::TensorStd,
) -> bool {
    conv3d_is_compatible(input, kernel, output)
}