wehub-resource-sync
138 lines
4.7 KiB
C++
138 lines
4.7 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 "paddle/phi/api/include/api.h"
|
|
|
|
namespace at {
|
|
|
|
// expand - expands tensor to new size
|
|
// PyTorch's expand works by right-aligning dimensions and broadcasting
|
|
// dimensions with size 1 to the target size
|
|
// Unlike Paddle's expand_v2, PyTorch allows non-singleton dimensions to be
|
|
// preserved when they match the corresponding target dimension
|
|
inline Tensor expand(const Tensor& self,
|
|
at::IntArrayRef size,
|
|
bool implicit = false) {
|
|
// implicit parameter is used by PyTorch's vmap for internal optimization.
|
|
// It doesn't affect the actual expand operation, so we can safely ignore it.
|
|
|
|
paddle::Tensor pd_tensor = self._PD_GetInner();
|
|
|
|
// Target sizes - convert to vector
|
|
std::vector<int64_t> target_size_vec(size.begin(), size.end());
|
|
auto target_rank = target_size_vec.size();
|
|
auto input_dims = pd_tensor.dims();
|
|
auto input_rank = static_cast<size_t>(input_dims.size());
|
|
|
|
// PyTorch's expand uses right-alignment semantics:
|
|
// - For 1D tensor expand to 2D: {3}.expand({3,4}) treats input as {3,1},
|
|
// expands to {3,4}
|
|
// - Non-singleton dimensions are preserved, singleton dimensions (1) can
|
|
// expand
|
|
//
|
|
// For example:
|
|
// {3}.expand({3, 4}) -> input {3} becomes {3, 1} implicitly
|
|
// then expand: dim 0: 3 stays 3, dim 1: 1 -> 4 -> result {3, 4}
|
|
|
|
if (input_rank < target_rank) {
|
|
// Add leading 1s to right-align with target shape (PyTorch behavior)
|
|
// Input {1, 2}, target {2, 3, 2} -> reshape to {1, 1, 2}
|
|
std::vector<int64_t> reshape_vec(target_rank, 1);
|
|
for (size_t i = 0; i < input_rank; ++i) {
|
|
reshape_vec[target_rank - input_rank + i] = input_dims[i];
|
|
}
|
|
|
|
// Check if Paddle's expand can handle this right-aligned shape
|
|
// Paddle allows: input[i] == 1 (can expand), or input[i] == target[i]
|
|
// (match)
|
|
bool can_use_paddle_expand = true;
|
|
size_t fail_dim = 0;
|
|
for (size_t i = 0; i < target_rank; ++i) {
|
|
bool dim_can_expand = (reshape_vec[i] == 1);
|
|
bool dim_is_matching = (reshape_vec[i] == target_size_vec[i]);
|
|
if (!dim_can_expand && !dim_is_matching) {
|
|
can_use_paddle_expand = false;
|
|
fail_dim = i;
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (can_use_paddle_expand) {
|
|
// Reshape to right-aligned shape, then expand
|
|
paddle::Tensor reshaped =
|
|
paddle::experimental::reshape(pd_tensor, phi::IntArray(reshape_vec));
|
|
paddle::Tensor result = paddle::experimental::expand(
|
|
reshaped, phi::IntArray(target_size_vec));
|
|
return Tensor(result);
|
|
}
|
|
|
|
PD_THROW("expand(): the expanded size of the tensor (",
|
|
target_size_vec[fail_dim],
|
|
") must match the existing size (",
|
|
reshape_vec[fail_dim],
|
|
") at non-singleton dimension ",
|
|
fail_dim,
|
|
".");
|
|
} else if (input_rank == target_rank) {
|
|
// Same rank - check if we can use expand directly
|
|
bool can_use_paddle_expand = true;
|
|
size_t fail_dim = 0;
|
|
for (size_t i = 0; i < target_rank; ++i) {
|
|
auto in_size = input_dims[i];
|
|
auto target_size = target_size_vec[i];
|
|
if (in_size != 1 && in_size != target_size) {
|
|
can_use_paddle_expand = false;
|
|
fail_dim = i;
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (can_use_paddle_expand) {
|
|
paddle::Tensor result = paddle::experimental::expand(
|
|
pd_tensor, phi::IntArray(target_size_vec));
|
|
return Tensor(result);
|
|
}
|
|
|
|
PD_THROW("expand(): the expanded size of the tensor (",
|
|
target_size_vec[fail_dim],
|
|
") must match the existing size (",
|
|
input_dims[fail_dim],
|
|
") at non-singleton dimension ",
|
|
fail_dim,
|
|
".");
|
|
} else {
|
|
PD_THROW("expand(): the number of sizes provided (",
|
|
target_rank,
|
|
") must be greater or equal to the number of dimensions in the "
|
|
"tensor (",
|
|
input_rank,
|
|
").");
|
|
}
|
|
}
|
|
|
|
} // namespace at
|
|
|
|
namespace at {
|
|
|
|
// Member function: Tensor::expand
|
|
inline Tensor Tensor::expand(at::IntArrayRef size, bool implicit) const {
|
|
return at::expand(*this, size, implicit);
|
|
}
|
|
|
|
} // namespace at
|