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) -> StaticName {
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        _node_id: usize,
50        state: &TurnState,
51        inputs: TVec<TValue>,
52    ) -> TractResult<TVec<TValue>> {
53        let input = args_1!(inputs);
54        if input.datum_type() == self.to {
55            Ok(tvec!(input))
56        } else if input.datum_type() == TDim::datum_type() {
57            let mut tmp = Tensor::zero_dt(i64::datum_type(), input.shape())?;
58            let input_dense = input.try_as_dense()?;
59            let mut tmp_dense = tmp.try_as_dense_mut()?;
60            for (dim, i) in tract_itertools::izip!(
61                input_dense.as_slice::<TDim>()?,
62                tmp_dense.as_slice_mut::<i64>()?
63            ) {
64                *i = dim.eval(&state.resolved_symbols).to_i64()?
65            }
66            Ok(tvec!(tmp.cast_to_dt(self.to)?.into_owned().into_tvalue()))
67        } else {
68            Ok(tvec!(input.cast_to_dt(self.to)?.into_owned().into_tvalue()))
69        }
70    }
71}
72
73impl TypedOp for Cast {
74    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
75        Ok(tvec!(self.to.fact(inputs[0].shape.clone())))
76    }
77
78    fn declutter(
79        &self,
80        model: &TypedModel,
81        node: &TypedNode,
82    ) -> TractResult<Option<TypedModelPatch>> {
83        if model.outlet_fact(node.inputs[0])?.datum_type == self.to {
84            TypedModelPatch::shunt_one_op(model, node)
85        } else {
86            Ok(None)
87        }
88    }
89
90    fn axes_mapping(
91        &self,
92        inputs: &[&TypedFact],
93        outputs: &[&TypedFact],
94    ) -> TractResult<AxesMapping> {
95        AxesMapping::natural(inputs, outputs)
96    }
97
98    fn change_axes(
99        &self,
100        model: &TypedModel,
101        node: &TypedNode,
102        _io: InOut,
103        change: &AxisOp,
104    ) -> TractResult<Option<AxisChangeConsequence>> {
105        Ok(Some(AxisChangeConsequence::new(model, node, None, change)))
106    }
107
108    fn slice(
109        &self,
110        patch: &mut TypedModelPatch,
111        _model: &TypedModel,
112        node: &TypedNode,
113        _prefix: &str,
114        inputs: &[OutletId],
115        _output_axis: usize,
116        _start: &TDim,
117        _end: &TDim,
118    ) -> TractResult<Option<TVec<OutletId>>> {
119        patch.wire_node(&node.name, &node.op, inputs).map(Some)
120    }
121
122    as_op!();
123}
tract_core/ops/cast.rs

tract_core/ops/
cast.rs
