| 2788 | public: |
| 2789 | SplitFunctor() {} |
| 2790 | Maybe<TensorTuple> operator()(const std::shared_ptr<one::Tensor>& x, |
| 2791 | const int64_t& split_size_or_sections, const int64_t& dim) const { |
| 2792 | int64_t axis = dim; |
| 2793 | axis = JUST(maybe_wrap_dim(axis, x->ndim())); |
| 2794 | CHECK_GE_OR_RETURN(split_size_or_sections, 0) |
| 2795 | << Error::RuntimeError() << "split expects split_size be non-negative, but got split_size=" |
| 2796 | << split_size_or_sections; |
| 2797 | int64_t dim_size = x->shape()->At(axis); |
| 2798 | int64_t num_splits = |
| 2799 | std::max<int64_t>((dim_size + split_size_or_sections - 1) / split_size_or_sections, 1); |
| 2800 | TensorTuple splits(num_splits); |
| 2801 | int64_t last_split_size = |
| 2802 | split_size_or_sections - (split_size_or_sections * num_splits - dim_size); |
| 2803 | for (int i = 0; i < num_splits; ++i) { |
| 2804 | int64_t length = i < num_splits - 1 ? split_size_or_sections : last_split_size; |
| 2805 | splits[i] = JUST(Narrow(x, axis, i * split_size_or_sections, length)); |
| 2806 | } |
| 2807 | return splits; |
| 2808 | } |
| 2809 | }; |
| 2810 | |
| 2811 | class UnbindFunctor { |
nothing calls this directly
no test coverage detected