| 72 | }; |
| 73 | |
| 74 | class ArrayPythonIterator { |
| 75 | public: |
| 76 | ArrayPythonIterator(mx::array x) : idx_(0), x_(std::move(x)) { |
| 77 | if (x_.shape(0) > 0 && x_.shape(0) < 10) { |
| 78 | splits_ = mx::split(x_, x_.shape(0)); |
| 79 | } |
| 80 | } |
| 81 | |
| 82 | mx::array next() { |
| 83 | if (idx_ >= x_.shape(0)) { |
| 84 | throw nb::stop_iteration(); |
| 85 | } |
| 86 | |
| 87 | if (idx_ >= 0 && idx_ < splits_.size()) { |
| 88 | return mx::squeeze(splits_[idx_++], 0); |
| 89 | } |
| 90 | |
| 91 | return *(x_.begin() + idx_++); |
| 92 | } |
| 93 | |
| 94 | private: |
| 95 | int idx_; |
| 96 | mx::array x_; |
| 97 | std::vector<mx::array> splits_; |
| 98 | }; |
| 99 | |
| 100 | void init_array(nb::module_& m) { |
| 101 | // Types |