1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

use crate::internal::*;
use crate::ops::element_wise::ElementWiseOp;

pub fn cast(to: DatumType) -> ElementWiseOp {
    ElementWiseOp(Box::new(Cast { to }))
}

#[derive(Debug, Clone, new, Hash)]
pub struct Cast {
    to: DatumType,
}

tract_linalg::impl_dyn_hash!(Cast);

impl ElementWiseMiniOp for Cast {
    fn name(&self) -> String {
        "Cast".into()
    }

    fn output_type(&self, _input_type: DatumType) -> Option<DatumType> {
        Some(self.to)
    }

    fn eval_out_of_place(&self, t: &Tensor) -> TractResult<Tensor> {
        t.cast_to_dt(self.to).map(|t| t.into_owned())
    }

    fn declutter(
        &self,
        model: &TypedModel,
        node: &TypedNode,
    ) -> TractResult<Option<TypedModelPatch>> {
        if model.outlet_fact(node.inputs[0])?.datum_type == self.to {
            Ok(Some(TypedModelPatch::shunt_one_op(model, node)?))
        } else {
            Ok(None)
        }
    }
}