| 3496 | } |
| 3497 | |
| 3498 | array scatter_axis( |
| 3499 | const array& a, |
| 3500 | const array& indices, |
| 3501 | const array& values, |
| 3502 | int axis, |
| 3503 | ScatterAxis::ReduceType mode, |
| 3504 | StreamOrDevice s) { |
| 3505 | std::string prefix = |
| 3506 | (mode == ScatterAxis::None) ? "[put_along_axis]" : "[scatter_add_axis]"; |
| 3507 | if (axis + a.ndim() < 0 || axis >= static_cast<int>(a.ndim())) { |
| 3508 | std::ostringstream msg; |
| 3509 | msg << prefix << " Received invalid axis for array with " << a.ndim() |
| 3510 | << " dimensions."; |
| 3511 | throw std::invalid_argument(msg.str()); |
| 3512 | } |
| 3513 | |
| 3514 | if (indices.ndim() != a.ndim()) { |
| 3515 | std::ostringstream msg; |
| 3516 | msg << prefix << " Indices of dimension " << indices.ndim() |
| 3517 | << " does not match array of dimension " << a.ndim() << "."; |
| 3518 | throw std::invalid_argument(msg.str()); |
| 3519 | } |
| 3520 | |
| 3521 | if (a.size() == 0) { |
| 3522 | return a; |
| 3523 | } |
| 3524 | |
| 3525 | auto upd = astype(values, a.dtype(), s); |
| 3526 | |
| 3527 | // Squeeze leading singletons out of update |
| 3528 | if (upd.ndim() > indices.ndim()) { |
| 3529 | std::vector<int> sq_ax(upd.ndim() - indices.ndim()); |
| 3530 | std::iota(sq_ax.begin(), sq_ax.end(), 0); |
| 3531 | upd = squeeze(upd, sq_ax, s); |
| 3532 | } |
| 3533 | |
| 3534 | auto inputs = broadcast_arrays({indices, upd}, s); |
| 3535 | inputs.insert(inputs.begin(), a); |
| 3536 | |
| 3537 | // Allow negative axis |
| 3538 | axis = axis < 0 ? a.ndim() + axis : axis; |
| 3539 | |
| 3540 | // Broadcast src, indices, values while ignoring the take axis |
| 3541 | inputs = broadcast_arrays(inputs, {axis - int(a.ndim())}, s); |
| 3542 | |
| 3543 | auto out_shape = inputs[0].shape(); |
| 3544 | return array( |
| 3545 | std::move(out_shape), |
| 3546 | a.dtype(), |
| 3547 | std::make_shared<ScatterAxis>(to_stream(s), mode, axis), |
| 3548 | std::move(inputs)); |
| 3549 | } |
| 3550 | |
| 3551 | array put_along_axis( |
| 3552 | const array& a, |