Skip to content

Commit 2d30104

Browse files
youge325claude
andcommitted
test(cholesky): align with Paddle compat cholesky behavior
新增 CholeskyTest 跨框架对比测试,覆盖: - Shape: 2x2, 3x3, 8x8, batch 2x3x3 - Dtype: kFloat, kDouble - API 变体: at::cholesky, Tensor::cholesky, upper=true/false - 异常: 非正定矩阵 更新 tensor_body.md 中 cholesky 状态为已对齐。 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent acefbad commit 2d30104

2 files changed

Lines changed: 270 additions & 1 deletion

File tree

doc/ATen/core/tensor_body.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@
7575
| `ccol_indices` ||| P2 | |
7676
| `ceil` ||| P2 | |
7777
| `chalf` ||| P2 | |
78-
| `cholesky` | | | P2 | |
78+
| `cholesky` | | | P2 | Tensor 方法与 at::cholesky 均已实现 |
7979
| `cholesky_inverse` ||| P2 | |
8080
| `cholesky_solve` ||| P2 | |
8181
| `clamp` ||| P2 | 已实现 Scalar 和 Tensor 版本 |

test/ATen/ops/CholeskyTest.cpp

Lines changed: 269 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,269 @@
1+
#include <ATen/ATen.h>
2+
#include <ATen/core/Tensor.h>
3+
#include <ATen/ops/cholesky.h>
4+
#include <gtest/gtest.h>
5+
6+
#include <cmath>
7+
#include <limits>
8+
#include <string>
9+
#include <vector>
10+
11+
#include "src/file_manager.h"
12+
13+
extern paddle_api_test::ThreadSafeParam g_custom_param;
14+
15+
namespace at {
16+
namespace test {
17+
18+
using paddle_api_test::FileManerger;
19+
using paddle_api_test::ThreadSafeParam;
20+
21+
// 使用 stride 按逻辑行优先顺序写出 tensor 元素,兼容不同内存布局
22+
static void write_cholesky_result_to_file(FileManerger* file,
23+
const at::Tensor& result) {
24+
*file << std::to_string(result.dim()) << " ";
25+
*file << std::to_string(result.numel()) << " ";
26+
for (int64_t i = 0; i < result.dim(); ++i) {
27+
*file << std::to_string(result.sizes()[i]) << " ";
28+
}
29+
30+
// 使用 stride 按逻辑顺序访问元素
31+
int64_t ndim = result.dim();
32+
std::vector<int64_t> strides(ndim);
33+
for (int64_t i = 0; i < ndim; ++i) {
34+
strides[i] = result.stride(i);
35+
}
36+
37+
switch (result.scalar_type()) {
38+
case at::kFloat: {
39+
float* data = result.data_ptr<float>();
40+
if (ndim == 2) {
41+
for (int64_t i = 0; i < result.size(0); ++i) {
42+
for (int64_t j = 0; j < result.size(1); ++j) {
43+
*file << std::to_string(data[i * strides[0] + j * strides[1]])
44+
<< " ";
45+
}
46+
}
47+
} else if (ndim == 3) {
48+
for (int64_t b = 0; b < result.size(0); ++b) {
49+
for (int64_t i = 0; i < result.size(1); ++i) {
50+
for (int64_t j = 0; j < result.size(2); ++j) {
51+
*file
52+
<< std::to_string(
53+
data[b * strides[0] + i * strides[1] + j * strides[2]])
54+
<< " ";
55+
}
56+
}
57+
}
58+
}
59+
break;
60+
}
61+
case at::kDouble: {
62+
double* data = result.data_ptr<double>();
63+
if (ndim == 2) {
64+
for (int64_t i = 0; i < result.size(0); ++i) {
65+
for (int64_t j = 0; j < result.size(1); ++j) {
66+
*file << std::to_string(data[i * strides[0] + j * strides[1]])
67+
<< " ";
68+
}
69+
}
70+
} else if (ndim == 3) {
71+
for (int64_t b = 0; b < result.size(0); ++b) {
72+
for (int64_t i = 0; i < result.size(1); ++i) {
73+
for (int64_t j = 0; j < result.size(2); ++j) {
74+
*file
75+
<< std::to_string(
76+
data[b * strides[0] + i * strides[1] + j * strides[2]])
77+
<< " ";
78+
}
79+
}
80+
}
81+
}
82+
break;
83+
}
84+
default: {
85+
*file << "unsupported_dtype ";
86+
break;
87+
}
88+
}
89+
}
90+
91+
// 构建对角占优的对称正定矩阵
92+
static at::Tensor make_spd_matrix(const std::vector<int64_t>& shape,
93+
at::ScalarType dtype) {
94+
at::Tensor A = at::zeros(shape, dtype);
95+
int64_t n = shape[shape.size() - 1];
96+
int64_t m = shape[shape.size() - 2];
97+
int64_t batch = 1;
98+
for (size_t i = 0; i + 2 < shape.size(); ++i) {
99+
batch *= shape[i];
100+
}
101+
102+
if (dtype == at::kFloat) {
103+
float* data = A.data_ptr<float>();
104+
for (int64_t b = 0; b < batch; ++b) {
105+
for (int64_t i = 0; i < m; ++i) {
106+
for (int64_t j = 0; j < n; ++j) {
107+
int64_t idx = b * m * n + i * n + j;
108+
if (i == j) {
109+
data[idx] = static_cast<float>(n);
110+
} else {
111+
data[idx] = 0.5f;
112+
}
113+
}
114+
}
115+
}
116+
} else if (dtype == at::kDouble) {
117+
double* data = A.data_ptr<double>();
118+
for (int64_t b = 0; b < batch; ++b) {
119+
for (int64_t i = 0; i < m; ++i) {
120+
for (int64_t j = 0; j < n; ++j) {
121+
int64_t idx = b * m * n + i * n + j;
122+
if (i == j) {
123+
data[idx] = static_cast<double>(n);
124+
} else {
125+
data[idx] = 0.5;
126+
}
127+
}
128+
}
129+
}
130+
}
131+
return A;
132+
}
133+
134+
class CholeskyTest : public ::testing::Test {
135+
protected:
136+
void SetUp() override {}
137+
};
138+
139+
// ========== 基础功能 ==========
140+
141+
TEST_F(CholeskyTest, BasicCholesky) {
142+
auto file_name = g_custom_param.get();
143+
FileManerger file(file_name);
144+
file.createFile();
145+
file << "BasicCholesky ";
146+
at::Tensor A = make_spd_matrix({3, 3}, at::kFloat);
147+
at::Tensor result = at::cholesky(A);
148+
write_cholesky_result_to_file(&file, result);
149+
file << "\n";
150+
file.saveFile();
151+
}
152+
153+
TEST_F(CholeskyTest, UpperTrue) {
154+
auto file_name = g_custom_param.get();
155+
FileManerger file(file_name);
156+
file.openAppend();
157+
file << "UpperTrue ";
158+
at::Tensor A = make_spd_matrix({3, 3}, at::kFloat);
159+
at::Tensor result = at::cholesky(A, /*upper=*/true);
160+
write_cholesky_result_to_file(&file, result);
161+
file << "\n";
162+
file.saveFile();
163+
}
164+
165+
// ========== Shape 覆盖 ==========
166+
167+
// 小矩阵 {2, 2}
168+
TEST_F(CholeskyTest, SmallMatrix) {
169+
auto file_name = g_custom_param.get();
170+
FileManerger file(file_name);
171+
file.openAppend();
172+
file << "SmallMatrix ";
173+
at::Tensor A = make_spd_matrix({2, 2}, at::kFloat);
174+
at::Tensor result = at::cholesky(A);
175+
write_cholesky_result_to_file(&file, result);
176+
file << "\n";
177+
file.saveFile();
178+
}
179+
180+
// 大矩阵 {8, 8}
181+
TEST_F(CholeskyTest, LargeMatrix) {
182+
auto file_name = g_custom_param.get();
183+
FileManerger file(file_name);
184+
file.openAppend();
185+
file << "LargeMatrix ";
186+
at::Tensor A = make_spd_matrix({8, 8}, at::kFloat);
187+
at::Tensor result = at::cholesky(A);
188+
write_cholesky_result_to_file(&file, result);
189+
file << "\n";
190+
file.saveFile();
191+
}
192+
193+
// batch 矩阵 {2, 3, 3}
194+
TEST_F(CholeskyTest, BatchMatrix) {
195+
auto file_name = g_custom_param.get();
196+
FileManerger file(file_name);
197+
file.openAppend();
198+
file << "BatchMatrix ";
199+
at::Tensor A = make_spd_matrix({2, 3, 3}, at::kFloat);
200+
at::Tensor result = at::cholesky(A);
201+
write_cholesky_result_to_file(&file, result);
202+
file << "\n";
203+
file.saveFile();
204+
}
205+
206+
// ========== Dtype 覆盖 ==========
207+
208+
// float64
209+
TEST_F(CholeskyTest, Float64Dtype) {
210+
auto file_name = g_custom_param.get();
211+
FileManerger file(file_name);
212+
file.openAppend();
213+
file << "Float64Dtype ";
214+
at::Tensor A = make_spd_matrix({3, 3}, at::kDouble);
215+
at::Tensor result = at::cholesky(A);
216+
write_cholesky_result_to_file(&file, result);
217+
file << "\n";
218+
file.saveFile();
219+
}
220+
221+
// ========== API 变体 ==========
222+
223+
// 方法调用 t.cholesky()
224+
TEST_F(CholeskyTest, MethodCholesky) {
225+
auto file_name = g_custom_param.get();
226+
FileManerger file(file_name);
227+
file.openAppend();
228+
file << "MethodCholesky ";
229+
at::Tensor A = make_spd_matrix({3, 3}, at::kFloat);
230+
at::Tensor result = A.cholesky();
231+
write_cholesky_result_to_file(&file, result);
232+
file << "\n";
233+
file.saveFile();
234+
}
235+
236+
// 方法调用 t.cholesky(true)
237+
TEST_F(CholeskyTest, MethodCholeskyUpper) {
238+
auto file_name = g_custom_param.get();
239+
FileManerger file(file_name);
240+
file.openAppend();
241+
file << "MethodCholeskyUpper ";
242+
at::Tensor A = make_spd_matrix({3, 3}, at::kFloat);
243+
at::Tensor result = A.cholesky(/*upper=*/true);
244+
write_cholesky_result_to_file(&file, result);
245+
file << "\n";
246+
file.saveFile();
247+
}
248+
249+
// ========== 异常测试 ==========
250+
251+
// 非正定矩阵应抛出异常
252+
TEST_F(CholeskyTest, NonPositiveDefinite) {
253+
auto file_name = g_custom_param.get();
254+
FileManerger file(file_name);
255+
file.openAppend();
256+
file << "NonPositiveDefinite ";
257+
at::Tensor A = at::zeros({3, 3}, at::kFloat);
258+
try {
259+
at::Tensor result = at::cholesky(A);
260+
write_cholesky_result_to_file(&file, result);
261+
} catch (const std::exception&) {
262+
file << "exception ";
263+
}
264+
file << "\n";
265+
file.saveFile();
266+
}
267+
268+
} // namespace test
269+
} // namespace at

0 commit comments

Comments
 (0)