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, PartialEq, Eq)]
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    op_out_of_plan!();
23
24    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
25        let input = args_1!(inputs);
26        let mut tensor = input.into_tensor();
27        unsafe {
28            tensor.set_shape_unchecked(&self.shape);
29        }
30        Ok(tvec!(tensor.into_tvalue()))
31    }
32}
33
34impl TypedOp for FiniteReshape {
35    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
36        Ok(tvec!(inputs[0].datum_type.fact(&self.shape)))
37    }
38
39    as_op!();
40}