tract_core/ops/array/
topk.rs

1use std::cmp::Ordering;
2
3use tract_data::itertools::Itertools;
4use tract_ndarray::{ArrayViewMutD, Axis, Dimension};
5
6use crate::internal::*;
7
8#[derive(Debug, Clone, new, Default, Hash, PartialEq, Eq)]
9pub struct Topk {
10    pub axis: usize,
11    pub largest: bool,
12    pub fallback_k: TDim,
13}
14
15impl Op for Topk {
16    fn mints_runtime_symbols(&self) -> bool {
17        true
18    }
19
20    fn name(&self) -> StaticName {
21        "Topk".into()
22    }
23
24    op_as_typed_op!();
25}
26
27impl EvalOp for Topk {
28    op_out_of_plan!();
29
30    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
31        let (input, k) = args_2!(inputs);
32        let mut output_shape: TVec<usize> = input.shape().into();
33        let k = k.cast_to_scalar::<i64>()? as usize;
34        output_shape[self.axis] = k;
35        let dt = input.datum_type();
36        let mut output_values = Tensor::zero_dt(dt, &output_shape)?;
37        let mut output_indices = Tensor::zero::<i64>(&output_shape)?;
38        let mut iterating_shape = output_shape.clone();
39        iterating_shape[self.axis] = 1;
40        let mut output_indices_plain = output_indices.try_as_plain_ram_mut()?;
41        let mut output_indices_view = output_indices_plain.to_array_view_mut::<i64>()?;
42        for coords in tract_ndarray::indices(&*iterating_shape) {
43            let mut coords: TVec<usize> = coords.as_array_view().as_slice().unwrap().into();
44            dispatch_numbers!(Self::inner_loop_t(dt)(
45                self,
46                &mut coords,
47                &input,
48                &mut output_values,
49                &mut output_indices_view,
50                k
51            ))?;
52        }
53        Ok(tvec!(output_values.into_tvalue(), output_indices.into_tvalue()))
54    }
55}
56
57impl Topk {
58    fn inner_loop_t<T: Datum + PartialOrd>(
59        &self,
60        coords: &mut [usize],
61        input: &Tensor,
62        output_values: &mut Tensor,
63        output_indices_view: &mut ArrayViewMutD<i64>,
64        k: usize,
65    ) -> TractResult<()> {
66        let mut output_values_plain = output_values.try_as_plain_ram_mut()?;
67        let mut output_values_view = output_values_plain.to_array_view_mut::<T>()?;
68        let mut view = input.to_plain_array_view::<T>()?;
69        for (ix, x) in coords.iter().enumerate() {
70            if ix != self.axis {
71                view.collapse_axis(Axis(ix), *x);
72            }
73        }
74        for (ix, (argmax, max)) in view
75            .iter()
76            .cloned()
77            .enumerate()
78            .sorted_by(|a, b| {
79                let ord = { a.1.partial_cmp(&b.1).unwrap_or(Ordering::Less) };
80                if self.largest { ord.reverse() } else { ord }
81            })
82            .take(k)
83            .enumerate()
84        {
85            coords[self.axis] = ix;
86            output_values_view[&*coords] = max;
87            output_indices_view[&*coords] = argmax as i64;
88        }
89        Ok(())
90    }
91}
92
93impl TypedOp for Topk {
94    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
95        let mut fact_values = inputs[0].without_value();
96        let mut fact_indices = inputs[0].without_value();
97        let k: TDim = if let Some(k) = &inputs[1].konst {
98            k.cast_to::<TDim>()?.try_as_plain_ram()?.to_scalar::<TDim>()?.clone()
99        } else {
100            self.fallback_k.clone()
101        };
102        fact_values.shape.set(self.axis, k.clone());
103        fact_indices.shape.set(self.axis, k);
104        fact_indices.datum_type = i64::datum_type();
105        Ok(tvec!(fact_values, fact_indices))
106    }
107
108    fn set_symbols(
109        &self,
110        _source: &TypedModel,
111        node: &TypedNode,
112        target: &mut TypedModel,
113        mapping: &HashMap<OutletId, OutletId>,
114        subs: &HashMap<Symbol, TDim>,
115    ) -> TractResult<TVec<OutletId>> {
116        let op = Topk {
117            axis: self.axis,
118            largest: self.largest,
119            fallback_k: self.fallback_k.substitute_all(subs)?,
120        };
121        let inputs = node.inputs.iter().map(|i| mapping[i]).collect::<TVec<_>>();
122        target.wire_node(&node.name, op, &inputs)
123    }
124
125    as_op!();
126}
tract_core/ops/array/topk.rs

tract_core/ops/array/
topk.rs
