From 31ece46f13e0e2ce0e39394be93d87edab550804 Mon Sep 17 00:00:00 2001 From: youge325 Date: Thu, 28 May 2026 22:23:39 +0800 Subject: [PATCH 1/7] test(take): align with Paddle compat take behavior Add cross-framework comparison tests for at::take: - Shape coverage: small 1D, multi-dim, scalar, empty, duplicate indices - Dtype coverage: kFloat, kDouble, kInt, kLong - Exception coverage: out-of-range index Related: PaddlePaddle/Paddle#79178 Co-Authored-By: Claude Opus 4.7 (1M context) --- test/ATen/ops/TakeTest.cpp | 227 +++++++++++++++++++++++++++++++++++++ 1 file changed, 227 insertions(+) create mode 100644 test/ATen/ops/TakeTest.cpp diff --git a/test/ATen/ops/TakeTest.cpp b/test/ATen/ops/TakeTest.cpp new file mode 100644 index 0000000..ffc2a25 --- /dev/null +++ b/test/ATen/ops/TakeTest.cpp @@ -0,0 +1,227 @@ +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "src/file_manager.h" + +extern paddle_api_test::ThreadSafeParam g_custom_param; + +namespace at { +namespace test { + +using paddle_api_test::FileManerger; + +class TakeTest : public ::testing::Test { + protected: + void SetUp() override {} +}; + +static void write_result_to_file(FileManerger* file, const at::Tensor& result) { + *file << std::to_string(result.dim()) << " "; + *file << std::to_string(result.numel()) << " "; + if (result.scalar_type() == at::kFloat) { + float* data = result.data_ptr(); + for (int64_t i = 0; i < result.numel(); ++i) { + *file << std::to_string(data[i]) << " "; + } + } else if (result.scalar_type() == at::kDouble) { + double* data = result.data_ptr(); + for (int64_t i = 0; i < result.numel(); ++i) { + *file << std::to_string(data[i]) << " "; + } + } else if (result.scalar_type() == at::kInt) { + int* data = result.data_ptr(); + for (int64_t i = 0; i < result.numel(); ++i) { + *file << std::to_string(data[i]) << " "; + } + } else if (result.scalar_type() == at::kLong) { + int64_t* data = result.data_ptr(); + for (int64_t i = 0; i < result.numel(); ++i) { + *file << std::to_string(data[i]) << " "; + } + } +} + +static at::Tensor make_test_tensor(at::ScalarType dtype) { + if (dtype == at::kFloat) { + auto t = at::zeros({3, 4}, at::kFloat); + float* data = t.data_ptr(); + for (int i = 0; i < 12; ++i) data[i] = static_cast(i); + return t; + } else if (dtype == at::kDouble) { + auto t = at::zeros({3, 4}, at::kDouble); + double* data = t.data_ptr(); + for (int i = 0; i < 12; ++i) data[i] = static_cast(i); + return t; + } else if (dtype == at::kInt) { + auto t = at::zeros({3, 4}, at::kInt); + int* data = t.data_ptr(); + for (int i = 0; i < 12; ++i) data[i] = i; + return t; + } else if (dtype == at::kLong) { + auto t = at::zeros({3, 4}, at::kLong); + int64_t* data = t.data_ptr(); + for (int i = 0; i < 12; ++i) data[i] = i; + return t; + } + return at::zeros({3, 4}, at::kFloat); +} + +static at::Tensor make_index_tensor(const std::vector& values) { + auto t = at::empty({static_cast(values.size())}, + at::TensorOptions().dtype(at::kLong)); + int64_t* data = t.data_ptr(); + for (size_t i = 0; i < values.size(); ++i) data[i] = values[i]; + return t; +} + +// Shape: small 1D index +TEST_F(TakeTest, TakeFloatSmall) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = make_index_tensor({0, 5, 11}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.createFile(); + file << "TakeFloatSmall "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeDoubleSmall) { + at::Tensor t = make_test_tensor(at::kDouble); + at::Tensor index = make_index_tensor({0, 5, 11}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeDoubleSmall "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeIntSmall) { + at::Tensor t = make_test_tensor(at::kInt); + at::Tensor index = make_index_tensor({0, 5, 11}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeIntSmall "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeLongSmall) { + at::Tensor t = make_test_tensor(at::kLong); + at::Tensor index = make_index_tensor({0, 5, 11}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeLongSmall "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +// Shape: multi-dimensional index +TEST_F(TakeTest, TakeFloatMultiDimIndex) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = at::empty({2, 2}, at::TensorOptions().dtype(at::kLong)); + int64_t* data = index.data_ptr(); + data[0] = 0; + data[1] = 3; + data[2] = 7; + data[3] = 10; + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeFloatMultiDimIndex "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +// Shape: scalar index (0-dim) +TEST_F(TakeTest, TakeFloatScalarIndex) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = at::empty({}, at::TensorOptions().dtype(at::kLong)); + index.data_ptr()[0] = 7; + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeFloatScalarIndex "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +// Shape: empty index (boundary) +TEST_F(TakeTest, TakeFloatEmptyIndex) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = at::empty({0}, at::TensorOptions().dtype(at::kLong)); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeFloatEmptyIndex "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +// Shape: duplicate indices +TEST_F(TakeTest, TakeFloatDuplicateIndices) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = make_index_tensor({1, 1, 3, 1}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeFloatDuplicateIndices "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +// Exception: index out of range +TEST_F(TakeTest, TakeExceptionOutOfRange) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = make_index_tensor({100}); // out of range + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeExceptionOutOfRange "; + try { + at::Tensor result = at::take(t, index); + write_result_to_file(&file, result); + } catch (const std::exception& e) { + file << "exception: "; + } + file << "\n"; + file.saveFile(); +} + +} // namespace test +} // namespace at From b83a5b1e18b54daf32b8d8d30e742cc9b7c18810 Mon Sep 17 00:00:00 2001 From: youge325 Date: Sun, 14 Jun 2026 16:39:56 +0800 Subject: [PATCH 2/7] test(take): cover index dtype and range alignment --- test/ATen/ops/TakeTest.cpp | 98 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/test/ATen/ops/TakeTest.cpp b/test/ATen/ops/TakeTest.cpp index ffc2a25..5ffe75d 100644 --- a/test/ATen/ops/TakeTest.cpp +++ b/test/ATen/ops/TakeTest.cpp @@ -81,6 +81,14 @@ static at::Tensor make_index_tensor(const std::vector& values) { return t; } +static at::Tensor make_int_index_tensor(const std::vector& values) { + auto t = at::empty({static_cast(values.size())}, + at::TensorOptions().dtype(at::kInt)); + int32_t* data = t.data_ptr(); + for (size_t i = 0; i < values.size(); ++i) data[i] = values[i]; + return t; +} + // Shape: small 1D index TEST_F(TakeTest, TakeFloatSmall) { at::Tensor t = make_test_tensor(at::kFloat); @@ -223,5 +231,95 @@ TEST_F(TakeTest, TakeExceptionOutOfRange) { file.saveFile(); } +TEST_F(TakeTest, TakeFloatNegativeIndex) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = make_index_tensor({-1, 0}); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeFloatNegativeIndex "; + try { + at::Tensor result = at::take(t, index); + write_result_to_file(&file, result); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeIntIndexThrows) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = make_int_index_tensor({0, 1}); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeIntIndexThrows "; + try { + at::Tensor result = at::take(t, index); + write_result_to_file(&file, result); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeExceptionAtNumel) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = make_index_tensor({t.numel()}); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeExceptionAtNumel "; + try { + at::Tensor result = at::take(t, index); + write_result_to_file(&file, result); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeExceptionBelowNegativeNumel) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = make_index_tensor({-t.numel() - 1}); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeExceptionBelowNegativeNumel "; + try { + at::Tensor result = at::take(t, index); + write_result_to_file(&file, result); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeCudaExceptionAtNumel) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = make_index_tensor({t.numel()}); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeCudaExceptionAtNumel "; + try { + at::Tensor result = at::take(t.cuda(), index.cuda()); + write_result_to_file(&file, result.cpu()); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + } // namespace test } // namespace at From e4db78055b78f898f812df888f2449728dd7643d Mon Sep 17 00:00:00 2001 From: youge325 Date: Sun, 14 Jun 2026 17:44:17 +0800 Subject: [PATCH 3/7] test(ATen): add CUDA take index alignment cases --- test/ATen/ops/TakeTest.cpp | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/test/ATen/ops/TakeTest.cpp b/test/ATen/ops/TakeTest.cpp index 5ffe75d..e446741 100644 --- a/test/ATen/ops/TakeTest.cpp +++ b/test/ATen/ops/TakeTest.cpp @@ -303,6 +303,24 @@ TEST_F(TakeTest, TakeExceptionBelowNegativeNumel) { file.saveFile(); } +TEST_F(TakeTest, TakeCudaNegativeIndex) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = make_index_tensor({-1}); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeCudaNegativeIndex "; + try { + at::Tensor result = at::take(t.cuda(), index.cuda()); + write_result_to_file(&file, result.cpu()); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + TEST_F(TakeTest, TakeCudaExceptionAtNumel) { at::Tensor t = make_test_tensor(at::kFloat); at::Tensor index = make_index_tensor({t.numel()}); @@ -321,5 +339,23 @@ TEST_F(TakeTest, TakeCudaExceptionAtNumel) { file.saveFile(); } +TEST_F(TakeTest, TakeCudaExceptionBelowNegativeNumel) { + at::Tensor t = make_test_tensor(at::kFloat); + at::Tensor index = make_index_tensor({-t.numel() - 1}); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeCudaExceptionBelowNegativeNumel "; + try { + at::Tensor result = at::take(t.cuda(), index.cuda()); + write_result_to_file(&file, result.cpu()); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + } // namespace test } // namespace at From 3b5d0d8645c78c6c0d000a8986a4e17a39c27675 Mon Sep 17 00:00:00 2001 From: youge325 Date: Mon, 15 Jun 2026 01:41:01 +0800 Subject: [PATCH 4/7] test(ATen): refine take index cases --- test/ATen/ops/TakeTest.cpp | 48 ++++++++++++++++++++++---------------- 1 file changed, 28 insertions(+), 20 deletions(-) diff --git a/test/ATen/ops/TakeTest.cpp b/test/ATen/ops/TakeTest.cpp index e446741..ecbf47c 100644 --- a/test/ATen/ops/TakeTest.cpp +++ b/test/ATen/ops/TakeTest.cpp @@ -1,6 +1,8 @@ #include #include +#include #include +#include #include #include #include @@ -74,19 +76,13 @@ static at::Tensor make_test_tensor(at::ScalarType dtype) { } static at::Tensor make_index_tensor(const std::vector& values) { - auto t = at::empty({static_cast(values.size())}, - at::TensorOptions().dtype(at::kLong)); - int64_t* data = t.data_ptr(); - for (size_t i = 0; i < values.size(); ++i) data[i] = values[i]; - return t; + return at::tensor(at::ArrayRef(values), + at::TensorOptions().dtype(at::kLong)); } static at::Tensor make_int_index_tensor(const std::vector& values) { - auto t = at::empty({static_cast(values.size())}, - at::TensorOptions().dtype(at::kInt)); - int32_t* data = t.data_ptr(); - for (size_t i = 0; i < values.size(); ++i) data[i] = values[i]; - return t; + return at::tensor(at::ArrayRef(values), + at::TensorOptions().dtype(at::kInt)); } // Shape: small 1D index @@ -149,12 +145,7 @@ TEST_F(TakeTest, TakeLongSmall) { // Shape: multi-dimensional index TEST_F(TakeTest, TakeFloatMultiDimIndex) { at::Tensor t = make_test_tensor(at::kFloat); - at::Tensor index = at::empty({2, 2}, at::TensorOptions().dtype(at::kLong)); - int64_t* data = index.data_ptr(); - data[0] = 0; - data[1] = 3; - data[2] = 7; - data[3] = 10; + at::Tensor index = make_index_tensor({0, 3, 7, 10}).reshape({2, 2}); at::Tensor result = at::take(t, index); auto file_name = g_custom_param.get(); @@ -169,8 +160,7 @@ TEST_F(TakeTest, TakeFloatMultiDimIndex) { // Shape: scalar index (0-dim) TEST_F(TakeTest, TakeFloatScalarIndex) { at::Tensor t = make_test_tensor(at::kFloat); - at::Tensor index = at::empty({}, at::TensorOptions().dtype(at::kLong)); - index.data_ptr()[0] = 7; + at::Tensor index = at::full({}, 7, at::kLong); at::Tensor result = at::take(t, index); auto file_name = g_custom_param.get(); @@ -197,6 +187,24 @@ TEST_F(TakeTest, TakeFloatEmptyIndex) { file.saveFile(); } +TEST_F(TakeTest, TakeEmptyInputNonEmptyIndexThrows) { + at::Tensor t = at::empty({0}, at::TensorOptions().dtype(at::kFloat)); + at::Tensor index = make_index_tensor({0}); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeEmptyInputNonEmptyIndexThrows "; + try { + at::Tensor result = at::take(t, index); + write_result_to_file(&file, result); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + // Shape: duplicate indices TEST_F(TakeTest, TakeFloatDuplicateIndices) { at::Tensor t = make_test_tensor(at::kFloat); @@ -224,8 +232,8 @@ TEST_F(TakeTest, TakeExceptionOutOfRange) { try { at::Tensor result = at::take(t, index); write_result_to_file(&file, result); - } catch (const std::exception& e) { - file << "exception: "; + } catch (const std::exception&) { + file << "exception "; } file << "\n"; file.saveFile(); From e147a6de1861c18de870c6f2ac91c781b6fb9318 Mon Sep 17 00:00:00 2001 From: youge325 Date: Mon, 29 Jun 2026 12:06:26 +0800 Subject: [PATCH 5/7] test(take): cover extended dtype alignment --- test/ATen/ops/TakeTest.cpp | 135 +++++++++++++++++++++++++++++++++++++ 1 file changed, 135 insertions(+) diff --git a/test/ATen/ops/TakeTest.cpp b/test/ATen/ops/TakeTest.cpp index ecbf47c..61575bd 100644 --- a/test/ATen/ops/TakeTest.cpp +++ b/test/ATen/ops/TakeTest.cpp @@ -5,8 +5,11 @@ #include #include #include +#include #include +#include +#include #include #include @@ -47,6 +50,40 @@ static void write_result_to_file(FileManerger* file, const at::Tensor& result) { for (int64_t i = 0; i < result.numel(); ++i) { *file << std::to_string(data[i]) << " "; } + } else if (result.scalar_type() == at::kBool) { + bool* data = result.data_ptr(); + for (int64_t i = 0; i < result.numel(); ++i) { + *file << std::to_string(static_cast(data[i])) << " "; + } + } else if (result.scalar_type() == at::kChar) { + int8_t* data = result.data_ptr(); + for (int64_t i = 0; i < result.numel(); ++i) { + *file << std::to_string(static_cast(data[i])) << " "; + } + } else if (result.scalar_type() == at::kHalf) { + at::Half* data = result.data_ptr(); + for (int64_t i = 0; i < result.numel(); ++i) { + *file << std::to_string(static_cast(data[i])) << " "; + } + } else if (result.scalar_type() == at::kBFloat16) { + at::BFloat16* data = result.data_ptr(); + for (int64_t i = 0; i < result.numel(); ++i) { + *file << std::to_string(static_cast(data[i])) << " "; + } + } else if (result.scalar_type() == at::kComplexFloat) { + c10::complex* data = result.data_ptr>(); + for (int64_t i = 0; i < result.numel(); ++i) { + auto value = static_cast>(data[i]); + *file << std::to_string(value.real()) << " "; + *file << std::to_string(value.imag()) << " "; + } + } else if (result.scalar_type() == at::kComplexDouble) { + c10::complex* data = result.data_ptr>(); + for (int64_t i = 0; i < result.numel(); ++i) { + auto value = static_cast>(data[i]); + *file << std::to_string(value.real()) << " "; + *file << std::to_string(value.imag()) << " "; + } } } @@ -80,6 +117,11 @@ static at::Tensor make_index_tensor(const std::vector& values) { at::TensorOptions().dtype(at::kLong)); } +static at::Tensor make_float_tensor(const std::vector& values) { + return at::tensor(at::ArrayRef(values), + at::TensorOptions().dtype(at::kFloat)); +} + static at::Tensor make_int_index_tensor(const std::vector& values) { return at::tensor(at::ArrayRef(values), at::TensorOptions().dtype(at::kInt)); @@ -142,6 +184,99 @@ TEST_F(TakeTest, TakeLongSmall) { file.saveFile(); } +TEST_F(TakeTest, TakeBoolSmall) { + at::Tensor t = make_index_tensor({0, 1, 0, 1}).to(at::kBool); + at::Tensor index = make_index_tensor({0, 3, 1}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeBoolSmall "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeCharSmall) { + std::vector values = {-4, -2, 3, 7}; + at::Tensor t = at::tensor(at::ArrayRef(values), + at::TensorOptions().dtype(at::kChar)); + at::Tensor index = make_index_tensor({0, 3, 1}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeCharSmall "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeHalfSmall) { + at::Tensor t = make_float_tensor({1.5f, -2.0f, 0.5f, 4.0f}).to(at::kHalf); + at::Tensor index = make_index_tensor({0, 3, 1}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeHalfSmall "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeBFloat16Small) { + at::Tensor t = + make_float_tensor({1.25f, -3.5f, 2.0f, 8.0f}).to(at::kBFloat16); + at::Tensor index = make_index_tensor({0, 3, 1}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeBFloat16Small "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeComplexFloatSmall) { + std::vector> values = { + {1.0f, 2.0f}, {3.0f, -4.0f}, {-5.0f, 6.0f}, {7.0f, 8.0f}}; + at::Tensor t = at::tensor(at::ArrayRef>(values), + at::TensorOptions().dtype(at::kComplexFloat)); + at::Tensor index = make_index_tensor({0, 3, 1}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeComplexFloatSmall "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeComplexDoubleSmall) { + std::vector> values = { + {1.0, -2.0}, {-3.0, 4.0}, {5.0, -6.0}, {-7.0, -8.0}}; + at::Tensor t = at::tensor(at::ArrayRef>(values), + at::TensorOptions().dtype(at::kComplexDouble)); + at::Tensor index = make_index_tensor({0, 3, 1}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeComplexDoubleSmall "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + // Shape: multi-dimensional index TEST_F(TakeTest, TakeFloatMultiDimIndex) { at::Tensor t = make_test_tensor(at::kFloat); From 4537d39d256bbb0c5e5360a2a891fa3b88a8d888 Mon Sep 17 00:00:00 2001 From: youge325 Date: Mon, 29 Jun 2026 14:10:00 +0800 Subject: [PATCH 6/7] test(take): cover strided tensor alignment --- test/ATen/ops/TakeTest.cpp | 76 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/test/ATen/ops/TakeTest.cpp b/test/ATen/ops/TakeTest.cpp index 61575bd..60aa126 100644 --- a/test/ATen/ops/TakeTest.cpp +++ b/test/ATen/ops/TakeTest.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -70,6 +71,13 @@ static void write_result_to_file(FileManerger* file, const at::Tensor& result) { for (int64_t i = 0; i < result.numel(); ++i) { *file << std::to_string(static_cast(data[i])) << " "; } + } else if (result.scalar_type() == at::kFloat8_e5m2 || + result.scalar_type() == at::kFloat8_e4m3fn) { + at::Tensor result_float = result.to(at::kFloat); + float* data = result_float.data_ptr(); + for (int64_t i = 0; i < result_float.numel(); ++i) { + *file << std::to_string(data[i]) << " "; + } } else if (result.scalar_type() == at::kComplexFloat) { c10::complex* data = result.data_ptr>(); for (int64_t i = 0; i < result.numel(); ++i) { @@ -243,6 +251,44 @@ TEST_F(TakeTest, TakeBFloat16Small) { file.saveFile(); } +TEST_F(TakeTest, TakeFloat8E5M2Small) { + at::Tensor t = + make_float_tensor({1.0f, 2.0f, 4.0f, 8.0f}).to(at::kFloat8_e5m2); + at::Tensor index = make_index_tensor({0, 3, 1}); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeFloat8E5M2Small "; + try { + at::Tensor result = at::take(t, index); + write_result_to_file(&file, result); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeFloat8E4M3FNSmall) { + at::Tensor t = + make_float_tensor({1.0f, 2.0f, 4.0f, 8.0f}).to(at::kFloat8_e4m3fn); + at::Tensor index = make_index_tensor({0, 3, 1}); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeFloat8E4M3FNSmall "; + try { + at::Tensor result = at::take(t, index); + write_result_to_file(&file, result); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + TEST_F(TakeTest, TakeComplexFloatSmall) { std::vector> values = { {1.0f, 2.0f}, {3.0f, -4.0f}, {-5.0f, 6.0f}, {7.0f, 8.0f}}; @@ -277,6 +323,36 @@ TEST_F(TakeTest, TakeComplexDoubleSmall) { file.saveFile(); } +TEST_F(TakeTest, TakeFloatNonContiguousInput) { + at::Tensor base = at::arange(6, at::TensorOptions().dtype(at::kFloat)); + at::Tensor t = base.as_strided({3}, {2}); + at::Tensor index = make_index_tensor({1}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeFloatNonContiguousInput "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeFloatNonContiguousIndex) { + at::Tensor t = at::arange(6, at::TensorOptions().dtype(at::kFloat)); + at::Tensor index_base = make_index_tensor({0, 1, 2, 3, 4, 5}); + at::Tensor index = index_base.as_strided({3}, {2}); + at::Tensor result = at::take(t, index); + + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeFloatNonContiguousIndex "; + write_result_to_file(&file, result); + file << "\n"; + file.saveFile(); +} + // Shape: multi-dimensional index TEST_F(TakeTest, TakeFloatMultiDimIndex) { at::Tensor t = make_test_tensor(at::kFloat); From 159ea40ecff7c7baf3aeec0cba034861b039dbb0 Mon Sep 17 00:00:00 2001 From: youge325 Date: Thu, 2 Jul 2026 09:21:38 +0800 Subject: [PATCH 7/7] test(take): cover unsigned dtype alignment --- test/ATen/ops/TakeTest.cpp | 78 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/test/ATen/ops/TakeTest.cpp b/test/ATen/ops/TakeTest.cpp index 60aa126..7e75fa1 100644 --- a/test/ATen/ops/TakeTest.cpp +++ b/test/ATen/ops/TakeTest.cpp @@ -51,6 +51,16 @@ static void write_result_to_file(FileManerger* file, const at::Tensor& result) { for (int64_t i = 0; i < result.numel(); ++i) { *file << std::to_string(data[i]) << " "; } + } else if (result.scalar_type() == at::kUInt16) { + uint16_t* data = result.data_ptr(); + for (int64_t i = 0; i < result.numel(); ++i) { + *file << std::to_string(data[i]) << " "; + } + } else if (result.scalar_type() == at::kUInt32) { + uint32_t* data = result.data_ptr(); + for (int64_t i = 0; i < result.numel(); ++i) { + *file << std::to_string(data[i]) << " "; + } } else if (result.scalar_type() == at::kBool) { bool* data = result.data_ptr(); for (int64_t i = 0; i < result.numel(); ++i) { @@ -251,6 +261,40 @@ TEST_F(TakeTest, TakeBFloat16Small) { file.saveFile(); } +TEST_F(TakeTest, TakeUInt16Throws) { + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeUInt16Throws "; + try { + at::Tensor t = make_index_tensor({2, 2, 2, 2}).to(at::kUInt16); + at::Tensor index = make_index_tensor({0, 3, 1}); + at::Tensor result = at::take(t, index); + write_result_to_file(&file, result); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeUInt32Throws) { + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeUInt32Throws "; + try { + at::Tensor t = make_index_tensor({2, 2, 2, 2}).to(at::kUInt32); + at::Tensor index = make_index_tensor({0, 3, 1}); + at::Tensor result = at::take(t, index); + write_result_to_file(&file, result); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + TEST_F(TakeTest, TakeFloat8E5M2Small) { at::Tensor t = make_float_tensor({1.0f, 2.0f, 4.0f, 8.0f}).to(at::kFloat8_e5m2); @@ -576,5 +620,39 @@ TEST_F(TakeTest, TakeCudaExceptionBelowNegativeNumel) { file.saveFile(); } +TEST_F(TakeTest, TakeCudaUInt16Throws) { + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeCudaUInt16Throws "; + try { + at::Tensor t = make_index_tensor({2, 2, 2, 2}).to(at::kUInt16); + at::Tensor index = make_index_tensor({0, 3, 1}); + at::Tensor result = at::take(t.cuda(), index.cuda()); + write_result_to_file(&file, result.cpu()); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + +TEST_F(TakeTest, TakeCudaUInt32Throws) { + auto file_name = g_custom_param.get(); + FileManerger file(file_name); + file.openAppend(); + file << "TakeCudaUInt32Throws "; + try { + at::Tensor t = make_index_tensor({2, 2, 2, 2}).to(at::kUInt32); + at::Tensor index = make_index_tensor({0, 3, 1}); + at::Tensor result = at::take(t.cuda(), index.cuda()); + write_result_to_file(&file, result.cpu()); + } catch (const std::exception&) { + file << "exception "; + } + file << "\n"; + file.saveFile(); +} + } // namespace test } // namespace at