| 508 | } |
| 509 | |
| 510 | std::tuple<std::vector<mx::array>, mx::array, std::vector<int>> |
| 511 | mlx_scatter_args_array( |
| 512 | const mx::array& src, |
| 513 | const mx::array& indices, |
| 514 | const mx::array& update) { |
| 515 | if (src.ndim() == 0) { |
| 516 | throw std::invalid_argument( |
| 517 | "too many indices for array: array is 0-dimensional"); |
| 518 | } |
| 519 | |
| 520 | auto up = squeeze_leading_singletons(update); |
| 521 | |
| 522 | // The update shape must broadcast with indices.shape + [1] + src.shape[1:] |
| 523 | auto up_shape = indices.shape(); |
| 524 | up_shape.insert(up_shape.end(), src.shape().begin() + 1, src.shape().end()); |
| 525 | up = broadcast_to(up, up_shape); |
| 526 | up_shape.insert(up_shape.begin() + indices.ndim(), 1); |
| 527 | up = reshape(up, up_shape); |
| 528 | |
| 529 | return {{indices}, up, {0}}; |
| 530 | } |
| 531 | |
| 532 | std::tuple<std::vector<mx::array>, mx::array, std::vector<int>> |
| 533 | mlx_scatter_args_slice( |
no test coverage detected