Skip to main content

tract_core/ops/
source.rs

1use crate::internal::*;
2
3#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
4pub struct TypedSource {
5    pub fact: TypedFact,
6}
7
8impl Op for TypedSource {
9    fn name(&self) -> StaticName {
10        "Source".into()
11    }
12    op_as_typed_op!();
13}
14
15impl EvalOp for TypedSource {
16    not_out_of_plan!();
17
18    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
19        ensure!(!inputs.is_empty(), "Input for node {} is missing", ctx.node_id);
20        Ok(inputs)
21    }
22}
23
24impl TypedOp for TypedSource {
25    fn output_facts(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
26        Ok(tvec!(self.fact.clone()))
27    }
28
29    fn change_axes(
30        &self,
31        model: &TypedModel,
32        node: &TypedNode,
33        _io: InOut,
34        change: &AxisOp,
35    ) -> TractResult<Option<AxisChangeConsequence>> {
36        let mut fact = self.fact.clone();
37        change.change_shape(&mut fact.shape, false)?;
38        Ok(Some(AxisChangeConsequence::new(
39            model,
40            node,
41            Some(Box::new(TypedSource::new(fact))),
42            change,
43        )))
44    }
45
46    fn set_symbols(
47        &self,
48        _source: &TypedModel,
49        node: &TypedNode,
50        target: &mut TypedModel,
51        _mapping: &HashMap<OutletId, OutletId>,
52        subs: &HashMap<Symbol, TDim>,
53    ) -> TractResult<TVec<OutletId>> {
54        let shape: TVec<_> =
55            self.fact.shape.iter().map(|d| d.substitute_all(subs)).collect::<TractResult<_>>()?;
56        target.wire_node(&node.name, Self { fact: self.fact.datum_type.fact(&*shape) }, &[])
57    }
58
59    as_op!();
60}