tract-core 0.2.0

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

use ops::prelude::*;
use std::fmt::Debug;

pub trait OpState: Debug {
    type Op: Op;
    fn eval(&mut self, op: &Self::Op, inputs: TVec<SharedTensor>) -> TractResult<TVec<SharedTensor>>;
}

impl OpState for Option<Box<OpState>> {
    fn eval(&mut self, op: &Op, inputs: TVec<SharedTensor>) -> TractResult<TVec<SharedTensor>> {
        match self {
            None => op.as_stateless().unwrap().eval(inputs),
            Some(state) => {
                let op = op.downcast_ref::<Self::Op>(op).unwrap();
                op.as_statefull().unwrap().eval(self, inputs)
            }
        }
    }
}

pub trait StatelessOp: Op {
    fn eval(&self, _inputs: TVec<SharedTensor>) -> TractResult<TVec<SharedTensor>>;
}

pub trait StatefullOp: Op {
    type State: OpState;
    fn dispatch_eval(&self, state: &mut OpState, inputs: TVec<SharedTensor>) -> TractResult<TVec<SharedTensor>> {
    }
}

pub trait OpStateDispatch {
    fn state(&self) -> TractResult<Option<Box<OpState>>>;
    fn as_stateless(&self) -> Option<&StatelessOp>;
}

impl<O: Op + StatelessOp + Clone> OpStateDispatch for O {
    fn state(&self) -> TractResult<Option<Box<OpState>>> {
        Ok(None)
    }

    fn as_stateless(&self) -> Option<&StatelessOp> {
        Some(self)
    }
}