tract_core/ops/
cast.rs

1use crate::internal::*;
2
3pub fn cast(to: DatumType) -> Cast {
4    Cast { to }
5}
6
7pub fn wire_cast(
8    prefix: impl AsRef<str>,
9    target: &mut TypedModel,
10    inputs: &[OutletId],
11    operating_datum_type: DatumType,
12) -> TractResult<TVec<OutletId>> {
13    let prefix = prefix.as_ref();
14    let mut wires = tvec!();
15    for mut wire in inputs.iter().copied() {
16        if target.outlet_fact(wire)?.datum_type != operating_datum_type {
17            wire = target.wire_node(
18                target.unique_name(format!("{prefix}.cast")),
19                crate::ops::cast::cast(operating_datum_type),
20                &[wire],
21            )?[0];
22        }
23        wires.push(wire);
24    }
25    Ok(wires)
26}
27
28#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
29pub struct Cast {
30    pub to: DatumType,
31}
32
33impl Op for Cast {
34    fn name(&self) -> Cow<str> {
35        "Cast".into()
36    }
37
38    op_as_typed_op!();
39    impl_op_same_as!();
40}
41
42impl EvalOp for Cast {
43    fn is_stateless(&self) -> bool {
44        true
45    }
46
47    fn eval_with_session(
48        &self,
49        state: &SessionState,
50        inputs: TVec<TValue>,
51    ) -> TractResult<TVec<TValue>> {
52        let input = args_1!(inputs);
53        if input.datum_type() == self.to {
54            Ok(tvec!(input))
55        } else if input.datum_type() == TDim::datum_type() {
56            let mut tmp = Tensor::zero_dt(i64::datum_type(), input.shape())?;
57            for (dim, i) in
58                tract_itertools::izip!(input.as_slice::<TDim>()?, tmp.as_slice_mut::<i64>()?)
59            {
60                *i = dim.eval(&state.resolved_symbols).to_i64()?
61            }
62            Ok(tvec!(tmp.cast_to_dt(self.to)?.into_owned().into_tvalue()))
63        } else {
64            Ok(tvec!(input.cast_to_dt(self.to)?.into_owned().into_tvalue()))
65        }
66    }
67}
68
69impl TypedOp for Cast {
70    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
71        Ok(tvec!(self.to.fact(inputs[0].shape.clone())))
72    }
73
74    fn declutter(
75        &self,
76        model: &TypedModel,
77        node: &TypedNode,
78    ) -> TractResult<Option<TypedModelPatch>> {
79        if model.outlet_fact(node.inputs[0])?.datum_type == self.to {
80            TypedModelPatch::shunt_one_op(model, node)
81        } else {
82            Ok(None)
83        }
84    }
85
86    fn axes_mapping(
87        &self,
88        inputs: &[&TypedFact],
89        outputs: &[&TypedFact],
90    ) -> TractResult<AxesMapping> {
91        AxesMapping::natural(inputs, outputs)
92    }
93
94    fn change_axes(
95        &self,
96        model: &TypedModel,
97        node: &TypedNode,
98        _io: InOut,
99        change: &AxisOp,
100    ) -> TractResult<Option<AxisChangeConsequence>> {
101        Ok(Some(AxisChangeConsequence::new(model, node, None, change)))
102    }
103
104    as_op!();
105}
tract_core/ops/cast.rs

tract_core/ops/
cast.rs
