tract-tensorflow 0.22.1

Tiny, no-nonsense, self contained, TensorFlow and ONNX inference
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

use tract_hir::internal::*;

use crate::model::ParsingContext;
use crate::tfpb::tensorflow::NodeDef;

use tract_hir::ops::array::GatherNd;

pub fn gather_nd(_ctx: &ParsingContext, pb: &NodeDef) -> TractResult<Box<dyn InferenceOp>> {
    let batch_dims = pb.get_attr_opt_int("batch_dims")?.unwrap_or(0);
    Ok(Box::new(GatherNd::new(batch_dims)))
}

#[cfg(test)]
mod tests {
    use super::*;

    // https://www.tensorflow.org/api_docs/python/tf/gather_nd
    #[test]
    fn simple_indexing() {
        let g = GatherNd::new(0);
        assert_eq!(
            g.eval(tvec!(tensor2(&[[1, 2], [3, 4]]).into(), tensor2(&[[0, 0], [1, 1]]).into()))
                .unwrap(),
            tvec!(tensor1(&[1, 4]).into())
        );
    }

    #[test]
    fn slice_indexing() {
        let g = GatherNd::new(0);
        assert_eq!(
            g.eval(tvec!(tensor2(&[[1, 2], [3, 4]]).into(), tensor2(&[[1], [0]]).into())).unwrap(),
            tvec!(tensor2(&[[3, 4], [1, 2]]).into())
        );
    }

    #[test]
    fn tensor_3d_1() {
        let g = GatherNd::new(0);
        let t = tensor3(&[[[10, 20], [30, 40]], [[11, 21], [31, 41]]]);
        assert_eq!(
            g.eval(tvec!(t.into(), tensor2(&[[1]]).into())).unwrap(),
            tvec!(tensor3(&[[[11, 21], [31, 41]]]).into())
        );
    }

    #[test]
    fn tensor_3d_2() {
        let g = GatherNd::new(0);
        let t = tensor3(&[[[10, 20], [30, 40]], [[11, 21], [31, 41]]]);
        assert_eq!(
            g.eval(tvec!(t.into(), tensor2(&[[0, 1], [1, 0]]).into())).unwrap(),
            tvec!(tensor2(&[[30, 40], [11, 21]]).into())
        );
    }

    #[test]
    fn tensor_3d_3() {
        let g = GatherNd::new(0);
        let t = tensor3(&[[[10, 20], [30, 40]], [[11, 21], [31, 41]]]);
        assert_eq!(
            g.eval(tvec!(t.into(), tensor2(&[[0, 0, 1], [1, 0, 1]]).into())).unwrap(),
            tvec!(tensor1(&[20, 21]).into())
        );
    }

    #[test]
    fn tensor_bd1_1() {
        let g = GatherNd::new(1);
        let t = tensor3(&[[[10, 20], [30, 40]], [[11, 21], [31, 41]]]);
        assert_eq!(
            g.eval(tvec!(t.into(), tensor2(&[[1], [0]]).into())).unwrap(),
            tvec!(tensor2(&[[30, 40], [11, 21]]).into())
        );
    }

    #[test]
    fn tensor_bd1_2() {
        let g = GatherNd::new(1);
        let t = tensor3(&[[[10, 20], [30, 40]], [[11, 21], [31, 41]]]);
        assert_eq!(
            g.eval(tvec!(t.into(), tensor3(&[[[1]], [[0]]]).into())).unwrap(),
            tvec!(tensor3(&[[[30, 40]], [[11, 21]]]).into())
        );
    }

    #[test]
    fn tensor_bd1_3() {
        let g = GatherNd::new(1);
        let t = tensor3(&[[[10, 20], [30, 40]], [[11, 21], [31, 41]]]);
        assert_eq!(
            g.eval(tvec!(t.into(), tensor3(&[[[1, 0]], [[0, 1]]]).into())).unwrap(),
            tvec!(tensor2(&[[30], [21]]).into())
        );
    }
}