| 2830 | public: |
| 2831 | ChunkFunctor() {} |
| 2832 | Maybe<TensorTuple> operator()(const std::shared_ptr<one::Tensor>& x, const int64_t& chunks, |
| 2833 | const int64_t& dim) const { |
| 2834 | const int64_t ndim = x->ndim(); |
| 2835 | int64_t infferd_dim = dim; |
| 2836 | CHECK_OR_RETURN(ndim > 0) << Error::RuntimeError() |
| 2837 | << "chunk expects at least a 1-dimensional tensor."; |
| 2838 | CHECK_OR_RETURN(chunks > 0) << Error::RuntimeError() |
| 2839 | << "chunk expects `chunks` to be greater than 0, got: " << chunks; |
| 2840 | infferd_dim = JUST(maybe_wrap_dim(infferd_dim, ndim)); |
| 2841 | |
| 2842 | const auto dim_size = x->shape()->At(infferd_dim); |
| 2843 | int64_t split_size = (dim_size + chunks - 1) / chunks; |
| 2844 | if (split_size == 0 && dim_size == 0) { |
| 2845 | std::vector<int64_t> split_sizes(chunks, split_size); |
| 2846 | split_sizes[chunks - 1] = split_size - (split_size * chunks - dim_size); |
| 2847 | return functional::SplitWithSize(x, split_sizes, infferd_dim); |
| 2848 | } else { |
| 2849 | return functional::Split(x, split_size, infferd_dim); |
| 2850 | } |
| 2851 | } |
| 2852 | }; |
| 2853 | |
| 2854 | class SplitLikeFunctor { |
nothing calls this directly
no test coverage detected