Chunk and combine function of argtopk Extract the indices of the k largest elements from a on the given axis. If k is negative, extract the indices of the -k smallest elements instead. Note that, unlike in the parent function, the returned elements are not sorted internally.
(a_plus_idx, k, axis, keepdims)
| 215 | |
| 216 | |
| 217 | def argtopk(a_plus_idx, k, axis, keepdims): |
| 218 | """Chunk and combine function of argtopk |
| 219 | |
| 220 | Extract the indices of the k largest elements from a on the given axis. |
| 221 | If k is negative, extract the indices of the -k smallest elements instead. |
| 222 | Note that, unlike in the parent function, the returned elements |
| 223 | are not sorted internally. |
| 224 | """ |
| 225 | assert keepdims is True |
| 226 | axis = axis[0] |
| 227 | |
| 228 | if isinstance(a_plus_idx, list): |
| 229 | a_plus_idx = list(flatten(a_plus_idx)) |
| 230 | a = np.concatenate([ai for ai, _ in a_plus_idx], axis) |
| 231 | idx = np.concatenate( |
| 232 | [np.broadcast_to(idxi, ai.shape) for ai, idxi in a_plus_idx], axis |
| 233 | ) |
| 234 | else: |
| 235 | a, idx = a_plus_idx |
| 236 | |
| 237 | if abs(k) >= a.shape[axis]: |
| 238 | return a_plus_idx |
| 239 | |
| 240 | idx2 = np.argpartition(a, -k, axis=axis) |
| 241 | k_slice = slice(-k, None) if k > 0 else slice(-k) |
| 242 | idx2 = idx2[tuple(k_slice if i == axis else slice(None) for i in range(a.ndim))] |
| 243 | return np.take_along_axis(a, idx2, axis), np.take_along_axis(idx, idx2, axis) |
| 244 | |
| 245 | |
| 246 | def argtopk_aggregate(a_plus_idx, k, axis, keepdims): |
no test coverage detected