wehub-resource-sync
64 lines
2.0 KiB
C++
64 lines
2.0 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>
|
|
|
|
namespace at {
|
|
|
|
inline at::Tensor unflatten(const at::Tensor& self,
|
|
int64_t dim,
|
|
at::IntArrayRef sizes) {
|
|
// Compute the new shape by replacing the dimension at 'dim' with 'sizes'
|
|
int64_t ndim = self._PD_GetInner().dims().size();
|
|
int64_t actual_dim = dim < 0 ? dim + ndim : dim;
|
|
std::vector<int64_t> new_shape;
|
|
for (int64_t i = 0; i < ndim; ++i) {
|
|
if (i == actual_dim) {
|
|
for (auto s : sizes) {
|
|
new_shape.push_back(s);
|
|
}
|
|
} else {
|
|
new_shape.push_back(self._PD_GetInner().dims()[i]);
|
|
}
|
|
}
|
|
return Tensor(paddle::experimental::reshape(self._PD_GetInner(), new_shape));
|
|
}
|
|
|
|
inline at::Tensor unflatten_symint(const at::Tensor& self,
|
|
int64_t dim,
|
|
c10::SymIntArrayRef sizes) {
|
|
return unflatten(
|
|
self,
|
|
dim,
|
|
at::IntArrayRef(reinterpret_cast<const int64_t*>(sizes.data()),
|
|
sizes.size()));
|
|
}
|
|
|
|
} // namespace at
|
|
|
|
namespace at {
|
|
|
|
inline at::Tensor Tensor::unflatten(int64_t dim, at::IntArrayRef sizes) const {
|
|
return at::unflatten(*this, dim, sizes);
|
|
}
|
|
|
|
inline at::Tensor Tensor::unflatten_symint(int64_t dim,
|
|
c10::SymIntArrayRef sizes) const {
|
|
return at::unflatten_symint(*this, dim, sizes);
|
|
}
|
|
|
|
} // namespace at
|