tract_tensorflow/ops/
mod.rs

1use tract_hir::internal::*;
2
3use crate::model::ParsingContext;
4use crate::model::TfOpRegister;
5use crate::tfpb::tensorflow::NodeDef;
6
7pub mod array;
8pub mod control_flow;
9pub mod logic;
10pub mod math;
11pub mod nn;
12pub mod quant;
13pub mod random;
14pub mod rec;
15// pub mod vars;
16
17pub fn register_all_ops(reg: &mut TfOpRegister) {
18    array::register_all_ops(reg);
19    control_flow::register_all_ops(reg);
20    logic::register_all_ops(reg);
21    math::register_all_ops(reg);
22    nn::register_all_ops(reg);
23    quant::register_all_ops(reg);
24    random::register_all_ops(reg);
25    rec::register_all_ops(reg);
26    // vars::register_all_ops(reg);
27    reg.insert("Cast", cast);
28    reg.insert("Const", konst);
29    reg.insert("Identity", |_, _| Ok(Box::new(tract_hir::ops::identity::Identity)));
30    reg.insert("NoOp", |_, _| Ok(Box::new(Noop)));
31    reg.insert("Placeholder", |_, _| Ok(Box::new(tract_hir::ops::source::Source::new())));
32}
33
34fn cast(_ctx: &ParsingContext, node: &NodeDef) -> TractResult<Box<dyn InferenceOp>> {
35    let dtype = node.get_attr_datum_type("DstT")?;
36    Ok(Box::new(tract_hir::ops::cast::cast(dtype)))
37}
38
39fn konst(_ctx: &ParsingContext, node: &NodeDef) -> TractResult<Box<dyn InferenceOp>> {
40    let dtype = node.get_attr_datum_type("dtype")?;
41    let mat = node.get_attr_tensor("value")?;
42
43    if mat.datum_type() != dtype {
44        bail!("Const node {:?} doesn't have the expected {:?} type.", mat, dtype);
45    }
46
47    Ok(Box::new(tract_hir::ops::konst::Const::new(mat.into())?))
48}
49
50#[derive(Clone, Debug, new, Hash, PartialEq, Eq)]
51pub struct Noop;
52
53impl Op for Noop {
54    fn name(&self) -> StaticName {
55        "Noop".into()
56    }
57
58    op_as_typed_op!();
59}
60
61impl EvalOp for Noop {
62    op_out_of_plan!();
63
64    fn eval(&self, _ctx: &EvalContext, _inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
65        Ok(tvec!(Tensor::from(false).into()))
66    }
67}
68
69impl InferenceRulesOp for Noop {
70    fn rules<'r, 'p: 'r, 's: 'r>(
71        &'s self,
72        s: &mut Solver<'r>,
73        _inputs: &'p [TensorProxy],
74        outputs: &'p [TensorProxy],
75    ) -> InferenceResult {
76        s.equals(&outputs[0].datum_type, bool::datum_type())?;
77        s.equals(&outputs[0].rank, 0)?;
78        Ok(())
79    }
80
81    as_op!();
82    to_typed!();
83}
84
85impl TypedOp for Noop {
86    fn output_facts(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
87        Ok(tvec!(bool::scalar_fact()))
88    }
89
90    as_op!();
91}
tract_tensorflow/ops/mod.rs

tract_tensorflow/ops/
mod.rs
