tract_core/ops/
konst.rs

1use crate::internal::*;
2
3#[derive(Debug, Clone, Hash, Eq, PartialEq)]
4pub struct Const(Arc<Tensor>, Option<Box<dyn OpaqueFact>>);
5
6impl Const {
7    pub fn new(tensor: Arc<Tensor>) -> TractResult<Const> {
8        Self::new_with_opt_opaque_fact(tensor, None)
9    }
10
11    pub fn new_with_opaque_fact(
12        tensor: Arc<Tensor>,
13        fact: Box<dyn OpaqueFact>,
14    ) -> TractResult<Const> {
15        Self::new_with_opt_opaque_fact(tensor, Some(fact))
16    }
17
18    pub fn new_with_opt_opaque_fact(
19        tensor: Arc<Tensor>,
20        fact: Option<Box<dyn OpaqueFact>>,
21    ) -> TractResult<Const> {
22        ensure!(fact.is_some() == tensor.datum_type().is_opaque());
23        Ok(Const(tensor, fact))
24    }
25
26    pub fn val(&self) -> &Arc<Tensor> {
27        &self.0
28    }
29
30    pub fn opaque_fact(&self) -> Option<&dyn OpaqueFact> {
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    impl_op_same_as!();
42}
43
44impl EvalOp for Const {
45    fn is_stateless(&self) -> bool {
46        true
47    }
48
49    fn eval(&self, _inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
50        Ok(tvec![Arc::clone(&self.0).into_tvalue()])
51    }
52}
53
54impl TypedOp for Const {
55    as_op!();
56
57    fn output_facts(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
58        let fact = TypedFact::from(&self.0);
59        if let Some(opaque) = &self.1 {
60            Ok(tvec!(fact.with_opaque_fact(opaque.clone())))
61        } else {
62            Ok(tvec!(fact))
63        }
64    }
65
66    fn cost(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
67        Ok(tvec!((Cost::Params(self.0.datum_type().unquantized()), self.0.len().into())))
68    }
69
70    fn concretize_dims(
71        &self,
72        _source: &TypedModel,
73        node: &TypedNode,
74        target: &mut TypedModel,
75        _mapping: &HashMap<OutletId, OutletId>,
76        values: &SymbolValues,
77    ) -> TractResult<TVec<OutletId>> {
78        let op = if self.0.datum_type() == TDim::datum_type() {
79            let mut tensor = self.0.clone().into_tensor();
80            for d in tensor.as_slice_mut::<TDim>()? {
81                *d = d.eval(values);
82            }
83            Const(tensor.into_arc_tensor(), self.1.clone())
84        } else {
85            self.clone()
86        };
87        target.wire_node(&node.name, op, &[])
88    }
89
90    fn change_axes(
91        &self,
92        _model: &TypedModel,
93        _node: &TypedNode,
94        io: InOut,
95        change: &AxisOp,
96    ) -> TractResult<Option<AxisChangeConsequence>> {
97        anyhow::ensure!(io == InOut::Out(0));
98        let mut new_tensor = self.0.clone().into_tensor();
99        if change.change_tensor(&mut new_tensor, false).is_ok() {
100            let mut sub = Const(new_tensor.into_arc_tensor(), None);
101            if self.1.is_some() {
102                let my_fact = self.output_facts(&[])?;
103                let changed_fact = change.output_facts(&[&my_fact[0]])?;
104                sub.1 = changed_fact[0].opaque_fact.clone();
105            }
106            Ok(Some(AxisChangeConsequence {
107                substitute_op: Some(Box::new(sub)),
108                wire_changes: tvec!((io, change.clone())),
109            }))
110        } else {
111            Ok(None)
112        }
113    }
114}
tract_core/ops/konst.rs

tract_core/ops/
konst.rs
