wehub-resource-sync
99 lines
2.8 KiB
C++
99 lines
2.8 KiB
C++
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#pragma once
|
|
|
|
#include <ATen/core/Tensor.h>
|
|
#include <vector>
|
|
|
|
#include "paddle/phi/api/include/api.h"
|
|
|
|
namespace at {
|
|
|
|
// chunk - splits tensor into chunks
|
|
inline std::vector<Tensor> chunk(const Tensor& self,
|
|
int64_t chunks,
|
|
int64_t dim = 0) {
|
|
if (chunks <= 0) {
|
|
PD_THROW("chunk expects chunks to be greater than 0, got ", chunks);
|
|
}
|
|
|
|
std::vector<Tensor> result;
|
|
paddle::Tensor pd_tensor = self._PD_GetInner();
|
|
|
|
int64_t rank = static_cast<int64_t>(pd_tensor.dims().size());
|
|
if (rank == 0) {
|
|
PD_THROW("chunk expects at least a 1-dimensional tensor");
|
|
}
|
|
|
|
int64_t original_dim = dim;
|
|
if (dim < 0) {
|
|
dim += rank;
|
|
}
|
|
if (dim < 0 || dim >= rank) {
|
|
PD_THROW("Dimension out of range (expected to be in range of [",
|
|
-rank,
|
|
", ",
|
|
rank - 1,
|
|
"], but got ",
|
|
original_dim,
|
|
")");
|
|
}
|
|
|
|
int64_t dim_size = pd_tensor.dims()[dim];
|
|
|
|
if (dim_size == 0) {
|
|
for (int64_t i = 0; i < chunks; ++i) {
|
|
auto chunk_tensor =
|
|
paddle::experimental::slice(pd_tensor, {dim}, {0}, {0}, {1}, {});
|
|
result.push_back(Tensor(chunk_tensor));
|
|
}
|
|
return result;
|
|
}
|
|
|
|
// PyTorch returns at most 'dim_size' non-empty chunks when chunks > dim_size
|
|
if (chunks > dim_size) {
|
|
chunks = dim_size;
|
|
}
|
|
|
|
int64_t chunk_size = (dim_size + chunks - 1) / chunks;
|
|
int64_t remaining = dim_size;
|
|
|
|
for (int64_t i = 0; i < chunks && remaining > 0; ++i) {
|
|
int64_t current_chunk_size = std::min(chunk_size, remaining);
|
|
auto chunk_tensor =
|
|
paddle::experimental::slice(pd_tensor,
|
|
{dim},
|
|
{i * chunk_size},
|
|
{i * chunk_size + current_chunk_size},
|
|
{1},
|
|
{});
|
|
result.push_back(Tensor(chunk_tensor));
|
|
remaining -= current_chunk_size;
|
|
}
|
|
|
|
return result;
|
|
}
|
|
|
|
} // namespace at
|
|
|
|
namespace at {
|
|
|
|
// Member function: Tensor::chunk
|
|
inline std::vector<Tensor> Tensor::chunk(int64_t chunks, int64_t dim) const {
|
|
return at::chunk(*this, chunks, dim);
|
|
}
|
|
|
|
} // namespace at
|