tract_tensorflow/ops/
control_flow.rs

1use tract_hir::internal::*;
2
3use crate::model::TfOpRegister;
4
5pub fn register_all_ops(reg: &mut TfOpRegister) {
6    reg.insert("Enter", |_, node| {
7        Ok(Box::new(LoopGate(LoopGateRole::Enter(node.get_attr_str("frame_name")?))))
8    });
9    reg.insert("Exit", |_, _| Ok(Box::new(LoopGate(LoopGateRole::Exit))));
10    reg.insert("LoopCond", |_, _| Ok(Box::new(LoopGate(LoopGateRole::LoopCond))));
11}
12
13#[derive(Debug, Clone, Hash, PartialEq, Eq)]
14pub enum LoopGateRole {
15    Enter(String),
16    Exit,
17    LoopCond,
18}
19
20#[derive(Debug, Clone, Hash, PartialEq, Eq)]
21pub struct LoopGate(LoopGateRole);
22
23impl Op for LoopGate {
24    fn name(&self) -> StaticName {
25        format!("{:?}", self.0).into()
26    }
27
28    not_a_typed_op!();
29}
30
31impl EvalOp for LoopGate {
32    op_out_of_plan!();
33
34    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
35        Ok(inputs)
36    }
37}
38
39impl InferenceRulesOp for LoopGate {
40    fn rules<'r, 'p: 'r, 's: 'r>(
41        &'s self,
42        s: &mut Solver<'r>,
43        inputs: &'p [TensorProxy],
44        outputs: &'p [TensorProxy],
45    ) -> InferenceResult {
46        check_input_arity(inputs, 1)?;
47        check_output_arity(outputs, 1)?;
48        s.equals(&inputs[0].datum_type, &outputs[0].datum_type)?;
49        s.equals(&inputs[0].shape, &outputs[0].shape)?;
50        Ok(())
51    }
52
53    as_op!();
54}
55
56#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
57pub enum NextIterationRole {
58    Source,
59    Sink,
60}
61
62#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
63pub struct NextIteration {
64    name: String,
65    role: NextIterationRole,
66}
67
68impl Op for NextIteration {
69    fn name(&self) -> StaticName {
70        format!("{:?}({})", self.role, self.name).into()
71    }
72
73    not_a_typed_op!();
74}
75
76impl EvalOp for NextIteration {
77    not_out_of_plan!();
78
79    fn state(&self, _ctx: &EvalContext) -> TractResult<Option<Box<dyn OpState>>> {
80        unimplemented!();
81    }
82}
83
84impl InferenceRulesOp for NextIteration {
85    fn rules<'r, 'p: 'r, 's: 'r>(
86        &'s self,
87        _s: &mut Solver<'r>,
88        inputs: &'p [TensorProxy],
89        outputs: &'p [TensorProxy],
90    ) -> InferenceResult {
91        match self.role {
92            NextIterationRole::Source => {
93                check_input_arity(inputs, 0)?;
94                check_output_arity(outputs, 1)?;
95            }
96            NextIterationRole::Sink => {
97                check_input_arity(inputs, 1)?;
98                check_output_arity(outputs, 0)?;
99            }
100        }
101        Ok(())
102    }
103
104    as_op!();
105}
tract_tensorflow/ops/control_flow.rs

tract_tensorflow/ops/
control_flow.rs
