Skip to main content

tract_core/ops/array/
reshape.rs

1use crate::internal::*;
2use tract_itertools::Itertools;
3
4#[derive(Debug, Clone, new, Default, Hash)]
5pub struct FiniteReshape {
6    pub shape: TVec<usize>,
7}
8
9impl Op for FiniteReshape {
10    fn name(&self) -> StaticName {
11        "Reshape".into()
12    }
13
14    fn info(&self) -> TractResult<Vec<String>> {
15        Ok(vec![format!("to shape: {}", self.shape.iter().join(","))])
16    }
17
18    op_as_typed_op!();
19}
20
21impl EvalOp for FiniteReshape {
22    fn is_stateless(&self) -> bool {
23        true
24    }
25
26    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
27        let input = args_1!(inputs);
28        let mut tensor = input.into_tensor();
29        unsafe {
30            tensor.set_shape_unchecked(&self.shape);
31        }
32        Ok(tvec!(tensor.into_tvalue()))
33    }
34}
35
36impl TypedOp for FiniteReshape {
37    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
38        Ok(tvec!(inputs[0].datum_type.fact(&self.shape)))
39    }
40
41    as_op!();
42}