| 3267 | } |
| 3268 | |
| 3269 | array gather( |
| 3270 | const array& a, |
| 3271 | const std::vector<array>& indices, |
| 3272 | const std::vector<int>& axes, |
| 3273 | const Shape& slice_sizes, |
| 3274 | StreamOrDevice s /* = {} */) { |
| 3275 | // Checks that indices, dimensions, and slice_sizes are all valid |
| 3276 | if (indices.size() > a.ndim()) { |
| 3277 | std::ostringstream msg; |
| 3278 | msg << "[gather] Too many index arrays. Got " << indices.size() |
| 3279 | << " index arrays for input with " << a.ndim() << " dimensions."; |
| 3280 | throw std::invalid_argument(msg.str()); |
| 3281 | } |
| 3282 | |
| 3283 | std::set dims(axes.begin(), axes.end()); |
| 3284 | if (dims.size() != axes.size()) { |
| 3285 | throw std::invalid_argument("[gather] Repeat axes not allowed in gather."); |
| 3286 | } |
| 3287 | if (!dims.empty() && (*dims.begin() < 0 || *dims.rbegin() >= a.ndim())) { |
| 3288 | throw std::invalid_argument("[gather] Axes don't match array dimensions."); |
| 3289 | } |
| 3290 | if (indices.size() != axes.size()) { |
| 3291 | throw std::invalid_argument( |
| 3292 | "[gather] Number of index arrays does not match number of axes."); |
| 3293 | } |
| 3294 | for (auto& x : indices) { |
| 3295 | if (x.dtype() == bool_) { |
| 3296 | throw std::invalid_argument("[Gather] Boolean indices not supported."); |
| 3297 | } |
| 3298 | } |
| 3299 | |
| 3300 | if (slice_sizes.size() != a.ndim()) { |
| 3301 | std::ostringstream msg; |
| 3302 | msg << "[gather] Got slice_sizes with size " << slice_sizes.size() |
| 3303 | << " for array with " << a.ndim() << " dimensions."; |
| 3304 | throw std::invalid_argument(msg.str()); |
| 3305 | } |
| 3306 | // Promote indices to the same type |
| 3307 | auto dtype = result_type(indices); |
| 3308 | if (issubdtype(dtype, inexact)) { |
| 3309 | throw std::invalid_argument( |
| 3310 | "[gather] Got indices with invalid dtype. Indices must be integral."); |
| 3311 | } |
| 3312 | |
| 3313 | // Broadcast and cast indices if necessary |
| 3314 | auto inputs = broadcast_arrays(indices); |
| 3315 | for (auto& idx : inputs) { |
| 3316 | idx = astype(idx, dtype, s); |
| 3317 | } |
| 3318 | |
| 3319 | if (a.size() == 0) { |
| 3320 | // Empty input, either the total slice size is 0 or the indices are empty |
| 3321 | auto total_slice = std::accumulate( |
| 3322 | slice_sizes.begin(), slice_sizes.end(), 1, std::multiplies<int64_t>{}); |
| 3323 | auto idx_size = !inputs.empty() ? inputs[0].size() : 1; |
| 3324 | if (idx_size != 0 && total_slice != 0) { |
| 3325 | std::ostringstream msg; |
| 3326 | msg << "[gather] If the input is empty, either the indices must be" |
no test coverage detected