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::*;
#[derive(Debug, Clone, new, Hash)]
pub struct Shape {
pub dt: DatumType,
}
tract_linalg::impl_dyn_hash!(Shape);
impl Op for Shape {
fn name(&self) -> Cow<str> {
"Shape".into()
}
op_core_mir!();
op_as_typed_op!();
not_a_pulsed_op!();
}
impl StatelessOp for Shape {
fn eval(&self, inputs: TVec<Arc<Tensor>>) -> TractResult<TVec<Arc<Tensor>>> {
let shape = rctensor1(&inputs[0].shape().iter().map(|&d| d as i64).collect::<Vec<_>>());
let shape = shape.cast_to_dt(self.dt)?.into_owned();
Ok(tvec![shape.into_arc_tensor()])
}
}
impl TypedOp for Shape {
fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
let shape = inputs[0].shape.iter().collect::<TVec<_>>();
let mut tensor = tensor1(&*shape);
if shape.iter().all(|d| d.to_integer().is_ok()) {
tensor = tensor.cast_to_dt(self.dt)?.into_owned();
}
Ok(tvec!(TypedFact::from(tensor)))
}
as_op!();
}