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 116 117 118 119
use tract_hir::internal::*; use crate::model::TfOpRegister; pub fn register_all_ops(reg: &mut TfOpRegister) { reg.insert("Enter", |_, node| { Ok(Box::new(LoopGate(LoopGateRole::Enter(node.get_attr_str("frame_name")?)))) }); reg.insert("Exit", |_, _| Ok(Box::new(LoopGate(LoopGateRole::Exit)))); reg.insert("LoopCond", |_, _| Ok(Box::new(LoopGate(LoopGateRole::LoopCond)))); } #[derive(Debug, Clone, Hash)] pub enum LoopGateRole { Enter(String), Exit, LoopCond, } #[derive(Debug, Clone, Hash)] pub struct LoopGate(LoopGateRole); impl_dyn_hash!(LoopGate); impl Op for LoopGate { fn name(&self) -> Cow<str> { format!("{:?}", self.0).into() } op_tf!(); not_a_typed_op!(); } impl EvalOp for LoopGate { fn is_stateless(&self) -> bool { true } fn eval(&self, inputs: TVec<Arc<Tensor>>) -> TractResult<TVec<Arc<Tensor>>> { Ok(inputs) } } impl InferenceRulesOp for LoopGate { fn rules<'r, 'p: 'r, 's: 'r>( &'s self, s: &mut Solver<'r>, inputs: &'p [TensorProxy], outputs: &'p [TensorProxy], ) -> InferenceResult { check_input_arity(&inputs, 1)?; check_output_arity(&outputs, 1)?; s.equals(&inputs[0].datum_type, &outputs[0].datum_type)?; s.equals(&inputs[0].shape, &outputs[0].shape)?; Ok(()) } as_op!(); } #[derive(Debug, Copy, Clone, PartialEq, Hash)] pub enum NextIterationRole { Source, Sink, } #[derive(Debug, Clone, new, Hash)] pub struct NextIteration { name: String, role: NextIterationRole, } impl_dyn_hash!(NextIteration); impl Op for NextIteration { fn name(&self) -> Cow<str> { format!("{:?}({})", self.role, self.name).into() } op_tf!(); not_a_typed_op!(); } impl EvalOp for NextIteration { fn is_stateless(&self) -> bool { false } fn state( &self, _state: &mut SessionState, _id: usize, ) -> TractResult<Option<Box<dyn OpState>>> { unimplemented!(); } } impl InferenceRulesOp for NextIteration { fn rules<'r, 'p: 'r, 's: 'r>( &'s self, _s: &mut Solver<'r>, inputs: &'p [TensorProxy], outputs: &'p [TensorProxy], ) -> InferenceResult { match self.role { NextIterationRole::Source => { check_input_arity(&inputs, 0)?; check_output_arity(&outputs, 1)?; } NextIterationRole::Sink => { check_input_arity(&inputs, 1)?; check_output_arity(&outputs, 0)?; } } Ok(()) } as_op!(); }