tract_core/ops/
konst.rs

1use crate::internal::*;
2
3#[derive(Debug, Clone, Hash, Eq, PartialEq)]
4pub struct Const(Arc<Tensor>, Option<Box<dyn ExoticFact>>);
5
6impl Const {
7    pub fn new(tensor: Arc<Tensor>) -> TractResult<Const> {
8        Self::new_with_opt_exotic_fact(tensor, None)
9    }
10
11    pub fn new_with_exotic_fact(
12        tensor: Arc<Tensor>,
13        fact: Box<dyn ExoticFact>,
14    ) -> TractResult<Const> {
15        Self::new_with_opt_exotic_fact(tensor, Some(fact))
16    }
17
18    pub fn new_with_opt_exotic_fact(
19        tensor: Arc<Tensor>,
20        fact: Option<Box<dyn ExoticFact>>,
21    ) -> TractResult<Const> {
22        ensure!(fact.is_some() || tensor.is_plain(), "Exotic tensor requires an exotic_fact");
23        Ok(Const(tensor, fact))
24    }
25
26    pub fn val(&self) -> &Arc<Tensor> {
27        &self.0
28    }
29
30    pub fn exotic_fact(&self) -> Option<&dyn ExoticFact> {
31        self.1.as_deref()
32    }
33}
34
35impl Op for Const {
36    fn name(&self) -> StaticName {
37        "Const".into()
38    }
39
40    op_as_typed_op!();
41}
42
43impl EvalOp for Const {
44    op_out_of_plan!();
45
46    fn eval(&self, _ctx: &EvalContext, _inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
47        Ok(tvec![Arc::clone(&self.0).into_tvalue()])
48    }
49}
50
51impl TypedOp for Const {
52    as_op!();
53
54    fn output_facts(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
55        let fact = if self.1.is_some() {
56            // Exotic const tensors (e.g. device-backed) may have storage that
57            // cannot produce an ExoticFact (like DeviceTensor). Build the fact
58            // from dt/shape and attach the explicit exotic_fact from self.1.
59            let mut f = TypedFact::dt_shape(
60                self.0.datum_type(),
61                ShapeFact::from_dims(self.0.shape().iter().map(TDim::from)),
62            );
63            f.konst = Some(Arc::clone(&self.0));
64            f.exotic_fact.clone_from(&self.1);
65            f
66        } else {
67            // Plain tensor: TryFrom sets uniform, uniform_tdim, exotic_fact from storage.
68            TypedFact::try_from(&self.0)?
69        };
70        Ok(tvec!(fact))
71    }
72
73    fn cost(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
74        Ok(tvec!((Cost::Params(self.0.datum_type().unquantized()), self.0.len().into())))
75    }
76
77    fn set_symbols(
78        &self,
79        _source: &TypedModel,
80        node: &TypedNode,
81        target: &mut TypedModel,
82        _mapping: &HashMap<OutletId, OutletId>,
83        subs: &HashMap<Symbol, TDim>,
84    ) -> TractResult<TVec<OutletId>> {
85        let op = if self.0.datum_type() == TDim::datum_type() {
86            let mut tensor = self.0.clone().into_tensor();
87            for d in tensor.try_as_plain_mut()?.as_slice_mut::<TDim>()? {
88                *d = d.substitute_all(subs)?;
89            }
90            Const(tensor.into_arc_tensor(), self.1.clone())
91        } else {
92            self.clone()
93        };
94        target.wire_node(&node.name, op, &[])
95    }
96
97    fn change_axes(
98        &self,
99        _model: &TypedModel,
100        _node: &TypedNode,
101        io: InOut,
102        change: &AxisOp,
103    ) -> TractResult<Option<AxisChangeConsequence>> {
104        anyhow::ensure!(io == InOut::Out(0));
105        let mut new_tensor = self.0.clone().into_tensor();
106        if change.change_tensor(&mut new_tensor, false).is_ok() {
107            let mut sub = Const(new_tensor.into_arc_tensor(), None);
108            if self.1.is_some() {
109                let my_fact = self.output_facts(&[])?;
110                let changed_fact = change.output_facts(&[&my_fact[0]])?;
111                sub.1 = changed_fact[0].exotic_fact.clone();
112            }
113            Ok(Some(AxisChangeConsequence {
114                substitute_op: Some(Box::new(sub)),
115                wire_changes: tvec!((io, change.clone())),
116            }))
117        } else {
118            Ok(None)
119        }
120    }
121}
tract_core/ops/konst.rs

tract_core/ops/
konst.rs
