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 mints_runtime_symbols(&self) -> bool {
10        true
11    }
12
13    fn name(&self) -> StaticName {
14        "Source".into()
15    }
16    op_as_typed_op!();
17}
18
19impl EvalOp for TypedSource {
20    not_out_of_plan!();
21
22    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
23        ensure!(!inputs.is_empty(), "Input for node {} is missing", ctx.node_id);
24        Ok(inputs)
25    }
26}
27
28impl TypedOp for TypedSource {
29    fn output_facts(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
30        Ok(tvec!(self.fact.clone()))
31    }
32
33    fn change_axes(
34        &self,
35        model: &TypedModel,
36        node: &TypedNode,
37        _io: InOut,
38        change: &AxisOp,
39    ) -> TractResult<Option<AxisChangeConsequence>> {
40        let mut fact = self.fact.clone();
41        // Block (rather than fail) on a change this source's shape cannot
42        // express: axis-change propagation reaching a scan body's input source
43        // with a wire-specific `Reshape` must stop the search, not error out.
44        if change.change_shape(&mut fact.shape, false).is_err() {
45            return Ok(None);
46        }
47        Ok(Some(AxisChangeConsequence::new(
48            model,
49            node,
50            Some(Box::new(TypedSource::new(fact))),
51            change,
52        )))
53    }
54
55    fn set_symbols(
56        &self,
57        _source: &TypedModel,
58        node: &TypedNode,
59        target: &mut TypedModel,
60        _mapping: &HashMap<OutletId, OutletId>,
61        subs: &HashMap<Symbol, TDim>,
62    ) -> TractResult<TVec<OutletId>> {
63        let shape: TVec<_> =
64            self.fact.shape.iter().map(|d| d.substitute_all(subs)).collect::<TractResult<_>>()?;
65        target.wire_node(&node.name, Self { fact: self.fact.datum_type.fact(&*shape) }, &[])
66    }
67
68    as_op!();
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    // A change this source's shape cannot express must block the axis-change
76    // search (Ok(None)) instead of erroring the whole optimization run. This
77    // is how a wire-specific Reshape change stops at a scan body's input
78    // source on the MMS TTS graph (#2928 follow-up diagnostics).
79    #[test]
80    fn blocks_inapplicable_change() -> TractResult<()> {
81        let mut model = TypedModel::default();
82        let s = model.add_source("s", f32::datum_type().fact([2usize, 3]))?;
83        let node = model.node(s.node);
84        let change =
85            AxisOp::Reshape(0, tvec![4.to_dim(), 5.to_dim()].into(), tvec![20.to_dim()].into());
86        let blocked = node
87            .op
88            .change_axes(&model, node, InOut::Out(0), &change)
89            .map(|r| r.is_none())
90            .unwrap_or(false);
91        assert!(blocked, "expected the source to block the change");
92        Ok(())
93    }
94}
tract_core/ops/source.rs

tract_core/ops/
source.rs
