tract_tensorflow/ops/array/
concatv2.rs

1use crate::model::ParsingContext;
2use crate::tfpb::tensorflow::NodeDef;
3use tract_hir::internal::*;
4use tract_hir::ops::array::TypedConcat;
5
6pub fn build(_ctx: &ParsingContext, _pb: &NodeDef) -> TractResult<Box<dyn InferenceOp>> {
7    Ok(expand(ConcatV2))
8}
9
10#[derive(Debug, Clone, new, Hash)]
11pub struct ConcatV2;
12
13
14
15impl Expansion for ConcatV2 {
16    fn name(&self) -> StaticName {
17        "ConcatV2".into()
18    }
19
20    fn rules<'r, 'p: 'r, 's: 'r>(
21        &'s self,
22        s: &mut Solver<'r>,
23        inputs: &'p [TensorProxy],
24        outputs: &'p [TensorProxy],
25    ) -> InferenceResult {
26        check_output_arity(outputs, 1)?;
27        let n = inputs.len() - 1;
28        s.equals_all((0..n).map(|i| (&inputs[i].datum_type).bex()).collect())?;
29        s.equals(&outputs[0].datum_type, &inputs[0].datum_type)?;
30        s.equals(&inputs[n].datum_type, DatumType::I32)?;
31        s.equals_all((0..n).map(|i| (&inputs[i].rank).bex()).collect())?;
32        s.equals(&inputs[n].rank, 0)?;
33        s.equals(&outputs[0].rank, &inputs[0].rank)?;
34        s.given(&inputs[n].value, move |s, axis| {
35            let axis = *axis.to_scalar::<i32>()? as usize;
36            trace!("axis for ConcatV2: {axis}");
37            for d in 0..axis {
38                s.equals_all((0..n).map(|i| (&inputs[i].shape[d]).bex()).collect())?;
39            }
40            for d in 0..axis {
41                s.equals(&inputs[0].shape[d], &outputs[0].shape[d])?;
42            }
43            s.given(&inputs[0].rank, move |s, rank| {
44                trace!("Given rank {rank}");
45                for d in (axis + 1)..(rank as usize) {
46                    s.equals(&inputs[0].shape[d], &outputs[0].shape[d])?;
47                }
48                for d in (axis + 1)..(rank as usize) {
49                    s.equals_all((0..n).map(|i| (&inputs[i].shape[d]).bex()).collect())?;
50                }
51                Ok(())
52            })?;
53
54            let mut concat_dim = -1 * outputs[0].shape[axis].bex();
55            for input in inputs.iter().take(n) {
56                concat_dim = concat_dim + input.shape[axis].bex();
57            }
58            s.equals_zero(concat_dim)
59        })
60    }
61
62    fn wire(
63        &self,
64        prefix: &str,
65        model: &mut TypedModel,
66        inputs: &[OutletId],
67    ) -> TractResult<TVec<OutletId>> {
68        if let Some(ref axis) = model.outlet_fact(*inputs.last().unwrap())?.konst {
69            let axis = *axis.to_scalar::<i32>()? as usize;
70            let inputs = inputs.iter().copied().rev().skip(1).rev().collect::<TVec<_>>();
71            model.wire_node(prefix, TypedConcat::new(axis), &inputs)
72        } else {
73            bail!("Except axis to be a constant")
74        }
75    }
76}
tract_tensorflow/ops/array/concatv2.rs

tract_tensorflow/ops/array/
concatv2.rs
