| 587 | } |
| 588 | |
| 589 | array expand_dims_impl( |
| 590 | const array& a, |
| 591 | std::vector<int> axes, |
| 592 | StreamOrDevice s /* = {} */) { |
| 593 | auto out_ndim = a.ndim() + axes.size(); |
| 594 | for (auto& ax : axes) { |
| 595 | auto new_ax = ax < 0 ? ax + out_ndim : ax; |
| 596 | if (new_ax < 0 || new_ax >= out_ndim) { |
| 597 | std::ostringstream msg; |
| 598 | msg << "[expand_dims] Invalid axis " << ax << " for output array with " |
| 599 | << a.ndim() << " dimensions."; |
| 600 | throw std::invalid_argument(msg.str()); |
| 601 | } |
| 602 | ax = new_ax; |
| 603 | } |
| 604 | auto shape = ExpandDims::output_shape(a, axes); |
| 605 | return array( |
| 606 | std::move(shape), |
| 607 | a.dtype(), |
| 608 | std::make_shared<ExpandDims>(to_stream(s), std::move(axes)), |
| 609 | {a}); |
| 610 | } |
| 611 | |
| 612 | array expand_dims(const array& a, int axis, StreamOrDevice s /* = {} */) { |
| 613 | return expand_dims_impl(a, {axis}, s); |
no test coverage detected