Scatter updates to given indices */
| 3568 | |
| 3569 | /** Scatter updates to given indices */ |
| 3570 | array scatter( |
| 3571 | const array& a, |
| 3572 | const std::vector<array>& indices, |
| 3573 | const array& updates, |
| 3574 | const std::vector<int>& axes, |
| 3575 | Scatter::ReduceType mode, |
| 3576 | StreamOrDevice s) { |
| 3577 | // Checks that indices, dimensions, and slice_sizes are all valid |
| 3578 | if (indices.size() > a.ndim()) { |
| 3579 | std::ostringstream msg; |
| 3580 | msg << "[scatter] Too many index arrays. Got " << indices.size() |
| 3581 | << " index arrays for input with " << a.ndim() << " dimensions."; |
| 3582 | throw std::invalid_argument(msg.str()); |
| 3583 | } |
| 3584 | for (auto& x : indices) { |
| 3585 | if (x.dtype() == bool_) { |
| 3586 | throw("[scatter] Boolean indices not supported."); |
| 3587 | } |
| 3588 | } |
| 3589 | |
| 3590 | std::set dims(axes.begin(), axes.end()); |
| 3591 | if (dims.size() != axes.size()) { |
| 3592 | throw std::invalid_argument( |
| 3593 | "[scatter] Repeat axes not allowed in scatter."); |
| 3594 | } |
| 3595 | if (!dims.empty() && (*dims.begin() < 0 || *dims.rbegin() >= a.ndim())) { |
| 3596 | throw std::invalid_argument("[scatter] Axes don't match array dimensions."); |
| 3597 | } |
| 3598 | if (indices.size() != axes.size()) { |
| 3599 | throw std::invalid_argument( |
| 3600 | "[scatter] Number of index arrays does not match number of axes."); |
| 3601 | } |
| 3602 | |
| 3603 | // Broadcast and cast indices if necessary |
| 3604 | auto inputs = broadcast_arrays(indices); |
| 3605 | |
| 3606 | Shape idx_shape; |
| 3607 | if (!inputs.empty()) { |
| 3608 | idx_shape = inputs[0].shape(); |
| 3609 | } |
| 3610 | |
| 3611 | if (updates.ndim() != (a.ndim() + idx_shape.size())) { |
| 3612 | std::ostringstream msg; |
| 3613 | msg << "[scatter] Updates with " << updates.ndim() |
| 3614 | << " dimensions does not match the sum of the array (" << a.ndim() |
| 3615 | << ") and indices (" << idx_shape.size() << ") dimensions."; |
| 3616 | throw std::invalid_argument(msg.str()); |
| 3617 | } |
| 3618 | for (int i = 0; i < idx_shape.size(); ++i) { |
| 3619 | if (updates.shape(i) != idx_shape[i]) { |
| 3620 | std::ostringstream msg; |
| 3621 | msg << "[scatter] Update shape " << updates.shape() |
| 3622 | << " is not valid for broadcasted index shape " << idx_shape << "."; |
| 3623 | throw std::invalid_argument(msg.str()); |
| 3624 | } |
| 3625 | } |
| 3626 | for (int i = 0; i < a.ndim(); ++i) { |
| 3627 | auto up_shape = updates.shape(i + idx_shape.size()); |
no test coverage detected