| 343 | } |
| 344 | |
| 345 | array::ArrayIterator::reference array::ArrayIterator::operator*() const { |
| 346 | auto start = Shape(arr.ndim(), 0); |
| 347 | auto end = arr.shape(); |
| 348 | auto shape = arr.shape(); |
| 349 | shape.erase(shape.begin()); |
| 350 | start[0] = idx; |
| 351 | end[0] = idx + 1; |
| 352 | return reshape(slice(arr, start, end), shape); |
| 353 | }; |
| 354 | |
| 355 | } // namespace mlx::core |