wehub-resource-sync
203 lines
7.1 KiB
C++
203 lines
7.1 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.
|
|
|
|
#include <ATen/Functions.h>
|
|
#include <ATen/core/TensorBody.h>
|
|
#include <ATen/cuda/EmptyTensor.h>
|
|
#include <ATen/native/cuda/Resize.h>
|
|
#include <ATen/ops/tensor.h>
|
|
#include <c10/core/ScalarType.h>
|
|
#include <c10/core/SymInt.h>
|
|
#include <c10/core/TensorOptions.h>
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
#include <c10/cuda/CUDAFunctions.h>
|
|
#include <c10/cuda/CUDAGuard.h>
|
|
#endif
|
|
#include "ATen/ATen.h"
|
|
#include "gtest/gtest.h"
|
|
#include "paddle/phi/common/float16.h"
|
|
#include "torch/all.h"
|
|
|
|
// ============================================================================
|
|
// Flatten Tests
|
|
// ============================================================================
|
|
|
|
TEST(TestFlatten, FlattenAllDims) {
|
|
// Test flatten with start_dim=0, end_dim=-1
|
|
// Flattens the entire tensor to 1D
|
|
at::Tensor tensor = at::ones({2, 3, 4}, at::kFloat);
|
|
at::Tensor flattened = tensor.flatten(0, -1);
|
|
|
|
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({24}));
|
|
ASSERT_EQ(flattened.numel(), tensor.numel());
|
|
}
|
|
|
|
TEST(TestFlatten, FlattenPartialDims) {
|
|
// Test flatten with specific start and end dimensions
|
|
at::Tensor tensor = at::ones({2, 3, 4, 5}, at::kFloat);
|
|
|
|
// Flatten dimensions 1 to 2 (3*4 = 12)
|
|
at::Tensor flattened = tensor.flatten(1, 2);
|
|
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({2, 12, 5}));
|
|
ASSERT_EQ(flattened.numel(), tensor.numel());
|
|
}
|
|
|
|
TEST(TestFlatten, FlattenSingleDim) {
|
|
// Test flatten when start_dim == end_dim (should be no-op)
|
|
at::Tensor tensor = at::ones({2, 3, 4}, at::kFloat);
|
|
|
|
at::Tensor flattened = tensor.flatten(1, 1);
|
|
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({2, 3, 4}));
|
|
}
|
|
|
|
TEST(TestFlatten, FlattenNegativeDims) {
|
|
// Test flatten with negative dimension indices
|
|
at::Tensor tensor = at::ones({2, 3, 4, 5}, at::kFloat);
|
|
|
|
// Flatten from -3 to -2 (dimensions 1 to 2)
|
|
at::Tensor flattened = tensor.flatten(-3, -2);
|
|
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({2, 12, 5}));
|
|
}
|
|
|
|
TEST(TestFlatten, FlattenFirstTwoDims) {
|
|
// Test flatten on first two dimensions
|
|
at::Tensor tensor = at::ones({2, 3, 4}, at::kFloat);
|
|
|
|
at::Tensor flattened = tensor.flatten(0, 1);
|
|
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({6, 4}));
|
|
}
|
|
|
|
TEST(TestFlatten, FlattenLastTwoDims) {
|
|
// Test flatten on last two dimensions
|
|
at::Tensor tensor = at::ones({2, 3, 4}, at::kFloat);
|
|
|
|
at::Tensor flattened = tensor.flatten(1, 2);
|
|
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({2, 12}));
|
|
}
|
|
|
|
TEST(TestFlatten, FlattenDataIntegrity) {
|
|
// Test that flatten preserves data
|
|
at::Tensor tensor = at::arange(24, at::kFloat).reshape({2, 3, 4});
|
|
at::Tensor flattened = tensor.flatten(0, -1);
|
|
|
|
const float* original_data = tensor.data_ptr<float>();
|
|
const float* flattened_data = flattened.data_ptr<float>();
|
|
|
|
for (int64_t i = 0; i < tensor.numel(); ++i) {
|
|
ASSERT_EQ(original_data[i], flattened_data[i]);
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Unflatten Tests
|
|
// ============================================================================
|
|
|
|
TEST(TestUnflatten, UnflattenBasic) {
|
|
// Test basic unflatten operation
|
|
at::Tensor tensor = at::ones({4, 6, 8}, at::kFloat);
|
|
|
|
// Unflatten dimension 1 (size 6) into (2, 3)
|
|
at::Tensor unflattened = tensor.unflatten(1, c10::IntArrayRef({2, 3}));
|
|
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({4, 2, 3, 8}));
|
|
ASSERT_EQ(unflattened.numel(), tensor.numel());
|
|
}
|
|
|
|
TEST(TestUnflatten, UnflattenFirstDim) {
|
|
// Test unflatten on first dimension
|
|
at::Tensor tensor = at::ones({6, 4}, at::kFloat);
|
|
|
|
// Unflatten dimension 0 (size 6) into (2, 3)
|
|
at::Tensor unflattened = tensor.unflatten(0, c10::IntArrayRef({2, 3}));
|
|
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({2, 3, 4}));
|
|
}
|
|
|
|
TEST(TestUnflatten, UnflattenLastDim) {
|
|
// Test unflatten on last dimension
|
|
at::Tensor tensor = at::ones({2, 12}, at::kFloat);
|
|
|
|
// Unflatten dimension 1 (size 12) into (3, 4)
|
|
at::Tensor unflattened = tensor.unflatten(1, c10::IntArrayRef({3, 4}));
|
|
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({2, 3, 4}));
|
|
}
|
|
|
|
TEST(TestUnflatten, UnflattenNegativeDim) {
|
|
// Test unflatten with negative dimension index
|
|
at::Tensor tensor = at::ones({4, 6, 8}, at::kFloat);
|
|
|
|
// Unflatten dimension -1 (last dim, size 8) into (4, 2)
|
|
at::Tensor unflattened = tensor.unflatten(-1, c10::IntArrayRef({4, 2}));
|
|
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({4, 6, 4, 2}));
|
|
}
|
|
|
|
TEST(TestUnflatten, UnflattenSymInt) {
|
|
// Test unflatten_symint (should behave same as unflatten)
|
|
at::Tensor tensor = at::ones({4, 6, 8}, at::kFloat);
|
|
|
|
// Unflatten dimension 1 using symint version
|
|
// Note: Must keep the underlying data alive
|
|
std::vector<c10::SymInt> sizes_vec = {2, 3};
|
|
c10::SymIntArrayRef sizes(sizes_vec);
|
|
at::Tensor unflattened = tensor.unflatten_symint(1, sizes);
|
|
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({4, 2, 3, 8}));
|
|
}
|
|
|
|
TEST(TestUnflatten, UnflattenDataIntegrity) {
|
|
// Test that unflatten preserves data
|
|
at::Tensor tensor = at::arange(24, at::kFloat).reshape({2, 12});
|
|
at::Tensor unflattened = tensor.unflatten(1, c10::IntArrayRef({3, 4}));
|
|
|
|
// Verify shape
|
|
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({2, 3, 4}));
|
|
|
|
// Verify numel
|
|
ASSERT_EQ(unflattened.numel(), tensor.numel());
|
|
}
|
|
|
|
// ============================================================================
|
|
// Flatten and Unflatten Combined Tests
|
|
// ============================================================================
|
|
|
|
TEST(TestFlattenUnflatten, RoundTrip) {
|
|
// Test that flatten followed by unflatten restores original shape
|
|
at::Tensor tensor = at::arange(24, at::kFloat).reshape({2, 3, 4});
|
|
|
|
// Flatten dimensions 1 and 2
|
|
at::Tensor flattened = tensor.flatten(1, 2);
|
|
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({2, 12}));
|
|
|
|
// Unflatten back to original shape
|
|
at::Tensor unflattened = flattened.unflatten(1, c10::IntArrayRef({3, 4}));
|
|
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({2, 3, 4}));
|
|
|
|
// Verify data integrity
|
|
ASSERT_EQ(tensor.numel(), unflattened.numel());
|
|
}
|
|
|
|
TEST(TestFlattenUnflatten, MultipleOperations) {
|
|
// Test multiple flatten/unflatten operations
|
|
at::Tensor tensor = at::ones({2, 3, 4, 5}, at::kFloat);
|
|
|
|
// Flatten all dimensions
|
|
at::Tensor flattened = tensor.flatten(0, -1);
|
|
ASSERT_EQ(flattened.sizes(), c10::IntArrayRef({120}));
|
|
|
|
// Unflatten into different shape
|
|
at::Tensor unflattened = flattened.unflatten(0, c10::IntArrayRef({6, 20}));
|
|
ASSERT_EQ(unflattened.sizes(), c10::IntArrayRef({6, 20}));
|
|
|
|
// Unflatten again
|
|
at::Tensor final_tensor = unflattened.unflatten(1, c10::IntArrayRef({4, 5}));
|
|
ASSERT_EQ(final_tensor.sizes(), c10::IntArrayRef({6, 4, 5}));
|
|
}
|