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        // A plain constant holds plain host bytes: storage that left them on a
24        // device brings them back rather than pinning the buffer for the life
25        // of the model.
26        let tensor = if fact.is_none() && !tensor.has_plain_ram_storage() {
27            Arc::new(Arc::unwrap_or_clone(tensor).into_plain_ram()?)
28        } else {
29            tensor
30        };
31        Ok(Const(tensor, fact))
32    }
33
34    pub fn val(&self) -> &Arc<Tensor> {
35        &self.0
36    }
37
38    pub fn exotic_fact(&self) -> Option<&dyn ExoticFact> {
39        self.1.as_deref()
40    }
41}
42
43impl Op for Const {
44    fn name(&self) -> StaticName {
45        "Const".into()
46    }
47
48    op_as_typed_op!();
49}
50
51impl EvalOp for Const {
52    op_out_of_plan!();
53
54    fn eval(&self, _ctx: &EvalContext, _inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
55        Ok(tvec![Arc::clone(&self.0).into_tvalue()])
56    }
57}
58
59impl TypedOp for Const {
60    as_op!();
61
62    fn output_facts(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
63        let fact = if self.1.is_some() {
64            // Exotic const tensors (e.g. device-backed) may have storage that
65            // cannot produce an ExoticFact (like DeviceTensor). Build the fact
66            // from dt/shape and attach the explicit exotic_fact from self.1.
67            let mut f = TypedFact::dt_shape(
68                self.0.datum_type(),
69                ShapeFact::from_dims(self.0.shape().iter().map(TDim::from)),
70            );
71            f.konst = Some(Arc::clone(&self.0));
72            f.exotic_fact.clone_from(&self.1);
73            f
74        } else {
75            // Plain tensor: TryFrom sets uniform, uniform_tdim, exotic_fact from storage.
76            TypedFact::try_from(&self.0)?
77        };
78        Ok(tvec!(fact))
79    }
80
81    fn cost(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
82        Ok(tvec!((Cost::Params(self.0.datum_type().unquantized()), self.0.len().into())))
83    }
84
85    fn set_symbols(
86        &self,
87        _source: &TypedModel,
88        node: &TypedNode,
89        target: &mut TypedModel,
90        _mapping: &HashMap<OutletId, OutletId>,
91        subs: &HashMap<Symbol, TDim>,
92    ) -> TractResult<TVec<OutletId>> {
93        let op = if self.0.datum_type() == TDim::datum_type() {
94            let mut tensor = self.0.clone().into_tensor();
95            for d in tensor.try_as_plain_ram_mut()?.as_slice_mut::<TDim>()? {
96                *d = d.substitute_all(subs)?;
97            }
98            Const(tensor.into_arc_tensor(), self.1.clone())
99        } else {
100            self.clone()
101        };
102        target.wire_node(&node.name, op, &[])
103    }
104
105    fn change_axes(
106        &self,
107        _model: &TypedModel,
108        _node: &TypedNode,
109        io: InOut,
110        change: &AxisOp,
111    ) -> TractResult<Option<AxisChangeConsequence>> {
112        anyhow::ensure!(io == InOut::Out(0));
113        let mut new_tensor = self.0.clone().into_tensor();
114        if change.change_tensor(&mut new_tensor, false).is_ok() {
115            let mut sub = Const(new_tensor.into_arc_tensor(), None);
116            if self.1.is_some() {
117                let my_fact = self.output_facts(&[])?;
118                let changed_fact = change.output_facts(&[&my_fact[0]])?;
119                sub.1 = changed_fact[0].exotic_fact.clone();
120            }
121            Ok(Some(AxisChangeConsequence {
122                substitute_op: Some(Box::new(sub)),
123                wire_changes: tvec!((io, change.clone())),
124            }))
125        } else {
126            Ok(None)
127        }
128    }
129}
tract_core/ops/konst.rs

tract_core/ops/
konst.rs
