| 542 | } |
| 543 | |
| 544 | std::shared_ptr<Array> sub( |
| 545 | const std::shared_ptr<const Array>& arr, |
| 546 | std::vector<int64_t> offset, |
| 547 | std::vector<int64_t> shape) { |
| 548 | if (arr->ndim() != offset.size()) { |
| 549 | throw std::runtime_error("Array: sub: array and offset dim mismatch"); |
| 550 | } |
| 551 | if (arr->ndim() != shape.size()) { |
| 552 | throw std::runtime_error("Array: sub: array and shape dim mismatch"); |
| 553 | } |
| 554 | |
| 555 | int64_t offset_sum = 0; |
| 556 | std::vector<int64_t> stride(arr->ndim(), 1); |
| 557 | for (int64_t dim = arr->ndim() - 1; dim >= 0; dim--) { |
| 558 | if (offset[dim] < 0 || offset[dim] >= arr->shape(dim)) { |
| 559 | throw std::runtime_error("Array: sub: offset out of bound"); |
| 560 | } |
| 561 | if (shape[dim] < 0) { |
| 562 | shape[dim] = arr->shape(dim); |
| 563 | } |
| 564 | if (offset[dim] + shape[dim] > arr->shape(dim)) { |
| 565 | throw std::runtime_error("Array: sub: shape out of bound"); |
| 566 | } |
| 567 | offset_sum += offset[dim] * stride[dim]; |
| 568 | if (dim > 0) { |
| 569 | stride[dim - 1] = stride[dim] * arr->shape(dim); |
| 570 | } |
| 571 | } |
| 572 | |
| 573 | auto res = std::make_shared<Array>(arr->type(), shape); |
| 574 | ARRAY_DISPATCH( |
| 575 | arr, |
| 576 | array_copy_strided_to_linear, |
| 577 | res->data(), |
| 578 | offset_sum, |
| 579 | arr->data(), |
| 580 | shape, |
| 581 | stride); |
| 582 | return res; |
| 583 | } |
| 584 | |
| 585 | } // namespace array |
| 586 | } // namespace data |