tract_hir/ops/array/
gather.rs

1use tract_core::ops::cast::cast;
2
3use crate::infer::*;
4use crate::internal::*;
5
6#[derive(Debug, Clone, new, Default, Hash)]
7pub struct Gather {
8    axis: i64,
9}
10
11impl Gather {
12    pub fn to_type_op(&self, input_rank: usize) -> tract_core::ops::array::Gather {
13        let axis = if self.axis < 0 { self.axis + input_rank as i64 } else { self.axis } as usize;
14        tract_core::ops::array::Gather::new(axis)
15    }
16}
17
18impl Expansion for Gather {
19    fn name(&self) -> StaticName {
20        "Gather".into()
21    }
22
23    fn wire(
24        &self,
25        prefix: &str,
26        model: &mut TypedModel,
27        inputs: &[OutletId],
28    ) -> TractResult<TVec<OutletId>> {
29        let input_rank = model.outlet_fact(inputs[0])?.rank();
30        let mut inputs: TVec<OutletId> = inputs.into();
31        inputs[1] = model.wire_node(
32            format!("{prefix}.cast_to_i64"),
33            cast(i64::datum_type()),
34            &[inputs[1]],
35        )?[0];
36        model.wire_node(prefix, self.to_type_op(input_rank), &inputs)
37    }
38
39    fn rules<'r, 'p: 'r, 's: 'r>(
40        &'s self,
41        s: &mut Solver<'r>,
42        inputs: &'p [TensorProxy],
43        outputs: &'p [TensorProxy],
44    ) -> InferenceResult {
45        check_input_arity(inputs, 2)?;
46        check_output_arity(outputs, 1)?;
47        s.equals(&inputs[0].datum_type, &outputs[0].datum_type)?;
48        s.equals(inputs[0].rank.bex() - 1 + inputs[1].rank.bex(), outputs[0].rank.bex())?;
49        s.given_2(&inputs[0].shape, &inputs[1].shape, move |s, input_shape, indices_shape| {
50            let rank = input_shape.len();
51            let output_shape =
52                self.to_type_op(rank).compute_output_shape(&input_shape, &indices_shape)?;
53            s.equals(&outputs[0].shape, output_shape)?;
54            Ok(())
55        })?;
56        Ok(())
57    }
58}
tract_hir/ops/array/gather.rs

tract_hir/ops/array/
gather.rs
