cheatah
Source BigBrain
Source

stdlib/tests/linalg_routines_test.cpp

1// Copyright (c) 2026 BigBrain LLC. MIT-licensed (see LICENSE).
2// Original work; see ACKNOWLEDGMENTS.md for the open-source ideas we build upon.
3#include "ndarray.hpp"
4#include "routines.hpp"
6#include <algorithm>
7#include <cmath>
8#include <complex>
9#include <stdexcept>
10#include <vector>
12#include <gtest/gtest.h>
14namespace nd = cheatah::ndarray;
15namespace la = cheatah::linalg;
17namespace {
18nd::NDArray mat(std::size_t r, std::size_t c, std::vector<double> data) {
19 return nd::reshape(nd::array(std::move(data)), {(long long)r, (long long)c});
21bool close(double a, double b, double tol = 1e-9) { return std::fabs(a - b) < tol; }
22// The general eig()/eigvals() return a complex spectrum; these read one element and
23// compare against a complex (or, implicitly, a real) expectation.
24std::complex<double> cget(const la::CNDArray& v, std::vector<long long> idx) {
25 return nd::get(v, std::move(idx));
27bool cclose(std::complex<double> a, std::complex<double> b, double tol = 1e-6) {
28 return std::abs(a - b) < tol;
30using C = std::complex<double>;
31la::CNDArray cvec(std::vector<C> data) { return nd::array(std::move(data)); }
32la::CNDArray cmat(std::size_t r, std::size_t c, std::vector<C> data) {
33 return nd::reshape(nd::array(std::move(data)), {(long long)r, (long long)c});
35} // namespace
37TEST(LinalgRoutines, ProductsAndTrace) {
38 const nd::NDArray a = mat(2, 3, {1, 2, 3, 4, 5, 6});
39 const nd::NDArray b = mat(3, 2, {7, 8, 9, 10, 11, 12});
40 const nd::NDArray c = la::matmul(a, b); // [[58,64],[139,154]]
41 EXPECT_DOUBLE_EQ(nd::get(c, {0, 0}), 58);
42 EXPECT_DOUBLE_EQ(nd::get(c, {0, 1}), 64);
43 EXPECT_DOUBLE_EQ(nd::get(c, {1, 1}), 154);
44 EXPECT_DOUBLE_EQ(la::dot(nd::array({1.0, 2.0, 3.0}), nd::array({4.0, 5.0, 6.0})), 32); // 4+10+18
45 EXPECT_DOUBLE_EQ(la::trace(mat(2, 2, {1, 2, 3, 4})), 5); // 1+4
48// The user-provided-output overload writes the SAME result into the caller's buffer with NO
49// reallocation: the buffer's data pointer is identical before and after, and the values are correct.
50TEST(LinalgRoutines, MatmulIntoReusesBuffer) {
51 const nd::NDArray a = mat(2, 3, {1, 2, 3, 4, 5, 6});
52 const nd::NDArray b = mat(3, 2, {7, 8, 9, 10, 11, 12});
53 nd::NDArray out = nd::zeros({2, 2});
54 const double* const before = out.buffer()->data() + out.offset(); // capture the buffer identity
55 la::matmul(out, a, b); // [[58,64],[139,154]] into out
56 EXPECT_EQ(out.buffer()->data() + out.offset(), before); // SAME buffer — no reallocation
57 EXPECT_DOUBLE_EQ(nd::get(out, {0, 0}), 58);
58 EXPECT_DOUBLE_EQ(nd::get(out, {0, 1}), 64);
59 EXPECT_DOUBLE_EQ(nd::get(out, {1, 0}), 139);
60 EXPECT_DOUBLE_EQ(nd::get(out, {1, 1}), 154);
61 // matches the allocating overload exactly
62 const nd::NDArray c = la::matmul(a, b);
63 EXPECT_DOUBLE_EQ(nd::get(out, {1, 1}), nd::get(c, {1, 1}));
64 // a wrong-shaped out, and an out that aliases an input (matmul is not in-place), are rejected.
65 nd::NDArray wrong = nd::zeros({3, 3});
66 EXPECT_THROW(la::matmul(wrong, a, b), std::runtime_error);
67 nd::NDArray sq = mat(2, 2, {1, 2, 3, 4});
68 EXPECT_THROW(la::matmul(sq, sq, sq), std::runtime_error);
69 // a non-2-D operand to the out-form is rejected up front ("expects 2-D matrices").
70 nd::NDArray vec = nd::array({1.0, 2.0, 3.0}); // 1-D
71 nd::NDArray out2 = nd::zeros({2, 2});
72 EXPECT_THROW(la::matmul(out2, vec, b), std::runtime_error);
75// Every product/least-squares front validates its operand dimensions and throws on a mismatch —
76// the error paths that a happy-path test never reaches. dot/vdot/inner reject unequal vector
77// lengths; batched matmul (two 3-D operands) rejects a mismatched contracted dimension; lstsq
78// rejects a row-count mismatch.
79TEST(LinalgRoutines, DimensionMismatchThrows) {
80 const nd::NDArray v2 = nd::array({1.0, 2.0});
81 const nd::NDArray v3 = nd::array({1.0, 2.0, 3.0});
82 EXPECT_THROW(la::dot(v2, v3), std::runtime_error);
83 EXPECT_THROW(la::vdot(v2, v3), std::runtime_error);
84 EXPECT_THROW(la::inner(v2, v3), std::runtime_error);
85 // Batched [B,M,K] @ [B,K,N]: equal batch counts but a mismatched inner dim (K vs K').
86 const nd::NDArray a3 = nd::reshape(nd::array(std::vector<double>(2 * 3 * 4, 1.0)), {2, 3, 4});
87 const nd::NDArray b3 = nd::reshape(nd::array(std::vector<double>(2 * 5 * 6, 1.0)), {2, 5, 6});
88 EXPECT_THROW(la::matmul(a3, b3), std::runtime_error); // K=4 != K'=5
89 // lstsq: A (m×n) and b (m×k) must share the row count m.
90 const nd::NDArray A = mat(3, 2, {1, 2, 3, 4, 5, 6});
91 const nd::NDArray rhs = mat(2, 1, {1, 2}); // 2 rows != A's 3
92 EXPECT_THROW(la::lstsq(A, rhs), std::runtime_error);
95// The memory-bound products / transposes have GENUINELY zero-allocation out-param overloads: they
96// write their kernel straight into the caller's buffer (data pointer identical before/after), match
97// the allocating overload, and reject a wrong shape or an out that aliases an input.
98TEST(LinalgRoutines, OuterIntoReusesBuffer) {
99 const nd::NDArray a = nd::array({1.0, 2.0, 3.0});
100 const nd::NDArray b = nd::array({4.0, 5.0});
101 nd::NDArray out = nd::zeros({3, 2});
102 const auto* before = out.buffer().get();
103 la::outer(out, a, b);
104 EXPECT_EQ(out.buffer().get(), before); // SAME buffer — no reallocation
105 const nd::NDArray ref = la::outer(a, b);
106 EXPECT_DOUBLE_EQ(nd::get(out, {0, 0}), 4.0); // 1*4
107 EXPECT_DOUBLE_EQ(nd::get(out, {2, 1}), nd::get(ref, {2, 1})); // 3*5, matches allocating form
108 nd::NDArray wrong = nd::zeros({2, 2});
109 EXPECT_THROW(la::outer(wrong, a, b), std::runtime_error);
112TEST(LinalgRoutines, KronIntoReusesBuffer) {
113 const nd::NDArray a = mat(2, 2, {1, 0, 0, 1});
114 const nd::NDArray b = mat(2, 2, {1, 2, 3, 4});
115 nd::NDArray out = nd::zeros({4, 4});
116 const auto* before = out.buffer().get();
117 la::kron(out, a, b);
118 EXPECT_EQ(out.buffer().get(), before);
119 const nd::NDArray ref = la::kron(a, b);
120 EXPECT_DOUBLE_EQ(nd::get(out, {0, 1}), nd::get(ref, {0, 1}));
121 EXPECT_DOUBLE_EQ(nd::get(out, {3, 3}), nd::get(ref, {3, 3}));
122 // A non-2-D operand is rejected up front ("kron expects 2-D matrices").
123 nd::NDArray vec = nd::array({1.0, 2.0}); // 1-D
124 EXPECT_THROW(la::kron(out, vec, b), std::runtime_error);
125 // An out that ALIASES an input is rejected by reject_alias (kron is not computed in place).
126 nd::NDArray alias = mat(2, 2, {1, 0, 0, 1});
127 EXPECT_THROW(la::kron(alias, alias, b), std::runtime_error);
130TEST(LinalgRoutines, ConjTransposeIntoReusesBuffer) {
131 const la::CNDArray M = cmat(2, 3, {C(1, 1), C(2, 0), C(3, -1), C(0, 2), C(1, 0), C(4, 4)});
132 la::CNDArray out = cmat(3, 2, std::vector<C>(6));
133 const auto* before = out.buffer().get();
134 la::conj_transpose(out, M);
135 EXPECT_EQ(out.buffer().get(), before);
136 const la::CNDArray ref = la::conj_transpose(M);
137 EXPECT_TRUE(cclose(cget(out, {0, 0}), C(1, -1)));
138 EXPECT_TRUE(cclose(cget(out, {2, 1}), cget(ref, {2, 1})));
141TEST(LinalgRoutines, ComplexMatmulIntoReusesBuffer) {
142 const la::CNDArray a = cmat(2, 3, {C(1, 0), C(2, 0), C(3, 0), C(4, 0), C(5, 0), C(6, 0)});
143 const la::CNDArray b = cmat(3, 2, {C(1, 1), C(0, 0), C(0, 1), C(1, 0), C(2, 0), C(0, 1)});
144 la::CNDArray out = cmat(2, 2, std::vector<C>(4));
145 const auto* before = out.buffer().get();
146 la::matmul(out, a, b);
147 EXPECT_EQ(out.buffer().get(), before);
148 const la::CNDArray ref = la::matmul(a, b);
149 EXPECT_TRUE(cclose(cget(out, {0, 0}), cget(ref, {0, 0})));
150 EXPECT_TRUE(cclose(cget(out, {1, 1}), cget(ref, {1, 1})));
153// The O(n³) factorizations reuse the caller's OUTPUT buffer (data pointer identical before/after)
154// and match the allocating overload — their internal factorization workspace is allocated regardless.
155TEST(LinalgRoutines, FactorizationOutReusesBuffer) {
156 const nd::NDArray A = mat(2, 2, {4, 3, 6, 3});
157 const nd::NDArray spd = mat(2, 2, {4, 2, 2, 3});
158 const nd::NDArray sym = mat(2, 2, {2, 1, 1, 2});
159 const nd::NDArray gen = mat(2, 2, {2, 0, 0, 5});
160 { // solve
161 nd::NDArray out = nd::zeros({2});
162 const auto* b = out.buffer().get();
163 la::solve(out, A, nd::array({10.0, 12.0}));
164 EXPECT_EQ(out.buffer().get(), b);
165 EXPECT_TRUE(close(nd::get(out, {0}), 1.0));
166 EXPECT_TRUE(close(nd::get(out, {1}), 2.0));
167 }
168 { // inv
169 nd::NDArray out = nd::zeros({2, 2});
170 const auto* b = out.buffer().get();
171 la::inv(out, A);
172 EXPECT_EQ(out.buffer().get(), b);
173 const nd::NDArray ref = la::inv(A);
174 EXPECT_TRUE(close(nd::get(out, {0, 0}), nd::get(ref, {0, 0})));
175 }
176 { // lstsq (2-D column rhs); on a square system it equals solve
177 nd::NDArray out = nd::zeros({2, 1});
178 const auto* b = out.buffer().get();
179 la::lstsq(out, A, mat(2, 1, {10, 12}));
180 EXPECT_EQ(out.buffer().get(), b);
181 EXPECT_TRUE(close(nd::get(out, {0, 0}), 1.0, 1e-6));
182 }
183 { // cholesky
184 nd::NDArray out = nd::zeros({2, 2});
185 const auto* b = out.buffer().get();
186 la::cholesky(out, spd);
187 EXPECT_EQ(out.buffer().get(), b);
188 const nd::NDArray ref = la::cholesky(spd);
189 EXPECT_TRUE(close(nd::get(out, {0, 0}), nd::get(ref, {0, 0})));
190 }
191 { // pinv
192 nd::NDArray out = nd::zeros({2, 2});
193 const auto* b = out.buffer().get();
194 la::pinv(out, A);
195 EXPECT_EQ(out.buffer().get(), b);
196 const nd::NDArray ref = la::pinv(A);
197 EXPECT_TRUE(close(nd::get(out, {0, 0}), nd::get(ref, {0, 0}), 1e-6));
198 }
199 { // matrix_power
200 nd::NDArray out = nd::zeros({2, 2});
201 const auto* b = out.buffer().get();
202 la::matrix_power(out, A, 2);
203 EXPECT_EQ(out.buffer().get(), b);
204 const nd::NDArray ref = la::matrix_power(A, 2);
205 EXPECT_TRUE(close(nd::get(out, {0, 0}), nd::get(ref, {0, 0})));
206 }
207 { // svdvals
208 nd::NDArray out = nd::zeros({2});
209 const auto* b = out.buffer().get();
210 la::svdvals(out, sym);
211 EXPECT_EQ(out.buffer().get(), b);
212 EXPECT_TRUE(close(nd::get(out, {0}), 3.0, 1e-6));
213 }
214 { // eigvalsh (symmetric)
215 nd::NDArray out = nd::zeros({2});
216 const auto* b = out.buffer().get();
217 la::eigvalsh(out, sym);
218 EXPECT_EQ(out.buffer().get(), b);
219 EXPECT_TRUE(close(nd::get(out, {0}), 3.0, 1e-6));
220 }
221 { // eigvals (general, complex out)
222 la::CNDArray out = cvec(std::vector<C>(2));
223 const auto* b = out.buffer().get();
224 la::eigvals(out, gen);
225 EXPECT_EQ(out.buffer().get(), b);
226 EXPECT_TRUE(cclose(cget(out, {0}), 5.0));
227 }
228 { // eigvalsh (complex Hermitian, real out)
229 const la::CNDArray H = cmat(2, 2, {C(2, 0), C(1, 1), C(1, -1), C(3, 0)});
230 nd::NDArray out = nd::zeros({2});
231 const auto* b = out.buffer().get();
232 la::eigvalsh(out, H);
233 EXPECT_EQ(out.buffer().get(), b);
234 EXPECT_TRUE(close(nd::get(out, {0}), 4.0, 1e-6));
235 }
238// The multi-output decompositions reuse EVERY caller-provided output buffer (one per factor).
239TEST(LinalgRoutines, DecompositionOutReusesBuffer) {
240 { // qr
241 const nd::NDArray A = mat(3, 2, {1, 0, 1, 1, 0, 1});
242 nd::NDArray q = nd::zeros({3, 2}), r = nd::zeros({2, 2});
243 const auto* bq = q.buffer().get();
244 const auto* br = r.buffer().get();
245 la::qr(q, r, A);
246 EXPECT_EQ(q.buffer().get(), bq);
247 EXPECT_EQ(r.buffer().get(), br);
248 const la::QR ref = la::qr(A);
249 EXPECT_TRUE(close(nd::get(q, {0, 0}), nd::get(ref.q, {0, 0}), 1e-6));
250 EXPECT_TRUE(close(nd::get(r, {0, 0}), nd::get(ref.r, {0, 0}), 1e-6));
251 }
252 { // svd
253 const nd::NDArray A = mat(2, 2, {2, 0, 0, 3});
254 nd::NDArray u = nd::zeros({2, 2}), s = nd::zeros({2}), vh = nd::zeros({2, 2});
255 const auto* bu = u.buffer().get();
256 const auto* bs = s.buffer().get();
257 const auto* bv = vh.buffer().get();
258 la::svd(u, s, vh, A);
259 EXPECT_EQ(u.buffer().get(), bu);
260 EXPECT_EQ(s.buffer().get(), bs);
261 EXPECT_EQ(vh.buffer().get(), bv);
262 EXPECT_TRUE(close(nd::get(s, {0}), 3.0, 1e-6));
263 }
264 { // eigh (symmetric, real)
265 nd::NDArray vals = nd::zeros({2}), vecs = nd::zeros({2, 2});
266 const auto* bvl = vals.buffer().get();
267 const auto* bvc = vecs.buffer().get();
268 la::eigh(vals, vecs, mat(2, 2, {2, 1, 1, 2}));
269 EXPECT_EQ(vals.buffer().get(), bvl);
270 EXPECT_EQ(vecs.buffer().get(), bvc);
271 EXPECT_TRUE(close(nd::get(vals, {0}), 3.0, 1e-6));
272 }
273 { // eigh (complex Hermitian: real values, complex vectors)
274 const la::CNDArray H = cmat(2, 2, {C(2, 0), C(1, 1), C(1, -1), C(3, 0)});
275 nd::NDArray vals = nd::zeros({2});
276 la::CNDArray vecs = cmat(2, 2, std::vector<C>(4));
277 const auto* bvl = vals.buffer().get();
278 const auto* bvc = vecs.buffer().get();
279 la::eigh(vals, vecs, H);
280 EXPECT_EQ(vals.buffer().get(), bvl);
281 EXPECT_EQ(vecs.buffer().get(), bvc);
282 EXPECT_TRUE(close(nd::get(vals, {0}), 4.0, 1e-6));
283 }
284 { // eig (general, complex values + vectors)
285 la::CNDArray vals = cvec(std::vector<C>(2));
286 la::CNDArray vecs = cmat(2, 2, std::vector<C>(4));
287 const auto* bvl = vals.buffer().get();
288 const auto* bvc = vecs.buffer().get();
289 la::eig(vals, vecs, mat(2, 2, {2, 0, 0, 5}));
290 EXPECT_EQ(vals.buffer().get(), bvl);
291 EXPECT_EQ(vecs.buffer().get(), bvc);
292 EXPECT_TRUE(cclose(cget(vals, {0}), 5.0));
293 }
296TEST(LinalgRoutines, SolveDetInv) {
297 const nd::NDArray A = mat(2, 2, {4, 3, 6, 3}); // det = 12-18 = -6
298 EXPECT_TRUE(close(la::det(A), -6.0));
299 const nd::NDArray x = la::solve(A, nd::array({10.0, 12.0})); // 4x+3y=10, 6x+3y=12 -> x=1,y=2
300 EXPECT_TRUE(close(nd::get(x, {0}), 1.0));
301 EXPECT_TRUE(close(nd::get(x, {1}), 2.0));
302 const nd::NDArray Ai = la::inv(A);
303 const nd::NDArray I = la::matmul(A, Ai); // identity
304 EXPECT_TRUE(close(nd::get(I, {0, 0}), 1.0));
305 EXPECT_TRUE(close(nd::get(I, {0, 1}), 0.0));
306 EXPECT_TRUE(close(nd::get(I, {1, 1}), 1.0));
309TEST(LinalgRoutines, CholeskyAndQR) {
310 const nd::NDArray A = mat(2, 2, {4, 2, 2, 3}); // SPD
311 const nd::NDArray L = la::cholesky(A);
312 const nd::NDArray LLt = la::matmul(L, la::matmul(la::inv(L), A)); // == A trivially; check L Lᵀ:
313 // verify L·Lᵀ == A
314 const nd::NDArray Lt = mat(2, 2, {nd::get(L, {0, 0}), nd::get(L, {1, 0}),
315 nd::get(L, {0, 1}), nd::get(L, {1, 1})});
316 const nd::NDArray rec = la::matmul(L, Lt);
317 EXPECT_TRUE(close(nd::get(rec, {0, 0}), 4.0));
318 EXPECT_TRUE(close(nd::get(rec, {1, 1}), 3.0));
320 const la::QR qr = la::qr(mat(3, 2, {1, 0, 1, 1, 0, 1}));
321 const nd::NDArray QtQ = la::matmul(la::pinv(qr.q), qr.q); // Q has orthonormal cols
322 EXPECT_TRUE(close(nd::get(QtQ, {0, 0}), 1.0, 1e-6));
323 const nd::NDArray reQR = la::matmul(qr.q, qr.r); // == original
324 EXPECT_TRUE(close(nd::get(reQR, {0, 0}), 1.0, 1e-6));
325 EXPECT_TRUE(close(nd::get(reQR, {1, 1}), 1.0, 1e-6));
328TEST(LinalgRoutines, SvdAndEigh) {
329 const nd::NDArray A = mat(2, 2, {2, 0, 0, 3});
330 const la::SVD s = la::svd(A);
331 EXPECT_TRUE(close(nd::get(s.s, {0}), 3.0, 1e-6)); // singular values 3, 2 (descending)
332 EXPECT_TRUE(close(nd::get(s.s, {1}), 2.0, 1e-6));
333 // svdvals (values-only fast path) agrees with svd().s
334 const nd::NDArray sv = la::svdvals(A);
335 EXPECT_TRUE(close(nd::get(sv, {0}), 3.0, 1e-6));
336 EXPECT_TRUE(close(nd::get(sv, {1}), 2.0, 1e-6));
338 // symmetric eigen: [[2,1],[1,2]] -> eigenvalues 3, 1
339 const la::Eig e = la::eigh(mat(2, 2, {2, 1, 1, 2}));
340 EXPECT_TRUE(close(nd::get(e.values, {0}), 3.0, 1e-6));
341 EXPECT_TRUE(close(nd::get(e.values, {1}), 1.0, 1e-6));
343 // general eigenvalues of [[2,0],[0,5]] -> 5, 2 (real, returned as complex)
344 const la::CNDArray ev = la::eigvals(mat(2, 2, {2, 0, 0, 5}));
345 EXPECT_TRUE(cclose(cget(ev, {0}), 5.0));
346 EXPECT_TRUE(cclose(cget(ev, {1}), 2.0));
349TEST(LinalgRoutines, ComplexProducts) {
350 const la::CNDArray a = cvec({C(1, 2), C(3, -1)});
351 const la::CNDArray b = cvec({C(0, 1), C(2, 0)});
352 // Bilinear dot (no conjugation): (1+2j)(0+1j) + (3-1j)(2) = (-2+1j) + (6-2j) = 4-1j.
353 EXPECT_TRUE(cclose(la::dot(a, b), C(4, -1)));
354 // Hermitian inner product (conjugate the first): conj(a)·b = (1-2j)(0+1j)+(3+1j)(2) = (2+1j)+(6+2j) = 8+3j.
355 EXPECT_TRUE(cclose(la::vdot(a, b), C(8, 3)));
356 // vdot(a,a) is the real squared norm ‖a‖² = 1+4+9+1 = 15.
357 EXPECT_TRUE(cclose(la::vdot(a, a), C(15, 0)));
359 // Conjugate transpose (Hermitian adjoint): transpose + conjugate every entry.
360 const la::CNDArray M = cmat(2, 2, {C(1, 1), C(2, 0), C(0, 0), C(3, -1)});
361 const la::CNDArray H = la::conj_transpose(M); // [[1-1j, 0],[2, 3+1j]]
362 EXPECT_TRUE(cclose(cget(H, {0, 0}), C(1, -1)));
363 EXPECT_TRUE(cclose(cget(H, {0, 1}), C(0, 0)));
364 EXPECT_TRUE(cclose(cget(H, {1, 0}), C(2, 0)));
365 EXPECT_TRUE(cclose(cget(H, {1, 1}), C(3, 1)));
367 // Complex matmul: M · Mᴴ is Hermitian; check entry (0,0) = |1+1j|² + |2|² = 2 + 4 = 6.
368 const la::CNDArray P = la::matmul(M, H);
369 EXPECT_TRUE(cclose(cget(P, {0, 0}), C(6, 0)));
370 EXPECT_TRUE(cclose(cget(P, {1, 1}), C(10, 0))); // |0|² + |3-1j|² = 0 + 10
372 // as_cvector accepts a 2-D N×1 / 1×N as a flat vector (like the real path)…
373 const la::CNDArray col = cmat(2, 1, {C(1, 0), C(0, 1)});
374 EXPECT_TRUE(cclose(la::dot(col, col), C(0, 0))); // 1·1 + i·i = 1 − 1 = 0
375 // …and rejects a genuine 2-D matrix where a vector is required.
376 EXPECT_THROW(la::vdot(M, M), std::runtime_error);
379TEST(LinalgRoutines, GeneralEigVectors) {
380 // For a real matrix, eig() returns complex eigenvalues AND eigenvectors (via
381 // inverse iteration). Verify A·v_k = λ_k·v_k for each column (phase-independent),
382 // building a complex copy Ac of A so we can multiply the complex eigenvectors.
383 const auto check = [](const std::vector<double>& data, std::size_t n) {
384 const nd::NDArray A = mat(n, n, data);
385 std::vector<C> cdata;
386 for (double x : data) cdata.push_back(C(x, 0.0));
387 const la::CNDArray Ac = cmat(n, n, cdata);
388 const la::EigC e = la::eig(A);
389 const la::CNDArray AV = la::matmul(Ac, e.vectors);
390 for (std::size_t k = 0; k < n; ++k) {
391 const C lam = cget(e.values, {(long long)k});
392 for (std::size_t r = 0; r < n; ++r) {
393 EXPECT_TRUE(cclose(cget(AV, {(long long)r, (long long)k}),
394 lam * cget(e.vectors, {(long long)r, (long long)k}), 1e-5));
395 }
396 }
397 };
398 check({2, 1, 0, 3}, 2); // real eigenvalues 3, 2 (upper-triangular)
399 check({0, -1, 1, 0}, 2); // complex conjugate pair ±i (rotation)
400 check({0, 1, 2, 0}, 2); // eigenvalues ±√2; forces a pivot row-swap in the solve
401 check({2, 1, 1, 1, 2, 1, 1, 1, 2}, 3); // symmetric -> real eigenvalues 4,1,1
404TEST(LinalgRoutines, ComplexHermitianEigh) {
405 // H = [[2, 1+i],[1-i, 3]] is Hermitian (conj_transpose(H) == H); eigenvalues 4, 1.
406 const la::CNDArray H = cmat(2, 2, {C(2, 0), C(1, 1), C(1, -1), C(3, 0)});
407 const nd::NDArray w = la::eigvalsh(H); // real, descending
408 EXPECT_TRUE(close(nd::get(w, {0}), 4.0, 1e-6));
409 EXPECT_TRUE(close(nd::get(w, {1}), 1.0, 1e-6));
411 const la::EighC e = la::eigh(H);
412 EXPECT_TRUE(close(nd::get(e.values, {0}), 4.0, 1e-6));
413 EXPECT_TRUE(close(nd::get(e.values, {1}), 1.0, 1e-6));
414 // Verify the eigenpairs: H·V should equal V·diag(λ), independent of eigenvector
415 // phase. So column k of H·V equals λ_k · column k of V.
416 const la::CNDArray HV = la::matmul(H, e.vectors);
417 for (int k = 0; k < 2; ++k) {
418 const double lam = nd::get(e.values, {k});
419 for (int r = 0; r < 2; ++r) {
420 EXPECT_TRUE(cclose(cget(HV, {r, k}), lam * cget(e.vectors, {r, k})));
421 }
422 }
423 // Eigenvectors are unit-norm: ⟨v,v⟩ = 1.
424 for (int k = 0; k < 2; ++k) {
425 const la::CNDArray vk = cvec({cget(e.vectors, {0, k}), cget(e.vectors, {1, k})});
426 EXPECT_TRUE(cclose(la::vdot(vk, vk), C(1, 0)));
427 }
430TEST(LinalgRoutines, NormAndRank) {
431 EXPECT_TRUE(close(la::norm(nd::array({3.0, 4.0})), 5.0)); // L2
432 EXPECT_EQ(la::matrix_rank(mat(2, 2, {1, 2, 2, 4})), 1); // rank-deficient
433 EXPECT_EQ(la::matrix_rank(mat(2, 2, {1, 0, 0, 1})), 2);
436TEST(LinalgRoutines, VdotInnerOuterKron) {
437 const nd::NDArray a = nd::array({1.0, 2.0, 3.0});
438 const nd::NDArray b = nd::array({4.0, 5.0, 6.0});
439 EXPECT_DOUBLE_EQ(la::vdot(a, b), 32.0); // 4+10+18
440 EXPECT_DOUBLE_EQ(la::inner(a, b), 32.0);
441 const nd::NDArray o = la::outer(nd::array({1.0, 2.0}), nd::array({3.0, 4.0})); // [[3,4],[6,8]]
442 EXPECT_DOUBLE_EQ(nd::get(o, {0, 0}), 3.0);
443 EXPECT_DOUBLE_EQ(nd::get(o, {1, 1}), 8.0);
444 const nd::NDArray k = la::kron(mat(2, 2, {1, 0, 0, 1}), mat(2, 2, {1, 2, 3, 4})); // I⊗B
445 EXPECT_DOUBLE_EQ(nd::get(k, {0, 0}), 1.0);
446 EXPECT_DOUBLE_EQ(nd::get(k, {0, 1}), 2.0);
447 EXPECT_DOUBLE_EQ(nd::get(k, {2, 2}), 1.0); // second diagonal block
448 EXPECT_DOUBLE_EQ(nd::get(k, {3, 3}), 4.0);
451TEST(LinalgRoutines, MatrixPower) {
452 const nd::NDArray A = mat(2, 2, {2, 0, 0, 3});
453 const nd::NDArray A0 = la::matrix_power(A, 0); // identity
454 EXPECT_TRUE(close(nd::get(A0, {0, 0}), 1.0));
455 EXPECT_TRUE(close(nd::get(A0, {1, 1}), 1.0));
456 const nd::NDArray A3 = la::matrix_power(A, 3); // diag(8, 27)
457 EXPECT_TRUE(close(nd::get(A3, {0, 0}), 8.0));
458 EXPECT_TRUE(close(nd::get(A3, {1, 1}), 27.0));
459 const nd::NDArray Am1 = la::matrix_power(A, -1); // diag(1/2, 1/3)
460 EXPECT_TRUE(close(nd::get(Am1, {0, 0}), 0.5));
461 EXPECT_TRUE(close(nd::get(Am1, {1, 1}), 1.0 / 3.0));
464TEST(LinalgRoutines, SlogdetAndCond) {
465 const la::SLogDet sd = la::slogdet(mat(2, 2, {4, 3, 6, 3})); // det = -6
466 EXPECT_TRUE(close(sd.sign, -1.0));
467 EXPECT_TRUE(close(sd.logabsdet, std::log(6.0), 1e-9));
468 EXPECT_TRUE(close(la::cond(mat(2, 2, {2, 0, 0, 2})), 1.0, 1e-6)); // well-conditioned
471TEST(LinalgRoutines, Lstsq) {
472 // lstsq(a, b) = pinv(a) · b and takes a 2-D b (column vector). On a square
473 // system, least-squares == solve.
474 const nd::NDArray A = mat(2, 2, {4, 3, 6, 3});
475 const nd::NDArray x = la::lstsq(A, mat(2, 1, {10, 12})); // -> [[1], [2]]
476 EXPECT_TRUE(close(nd::get(x, {0, 0}), 1.0, 1e-6));
477 EXPECT_TRUE(close(nd::get(x, {1, 0}), 2.0, 1e-6));
480TEST(LinalgRoutines, EigvalshSymmetric) {
481 const nd::NDArray w = la::eigvalsh(mat(2, 2, {2, 1, 1, 2})); // eigenvalues 1, 3
482 const double v0 = nd::get(w, {0}), v1 = nd::get(w, {1});
483 EXPECT_TRUE(close(std::min(v0, v1), 1.0, 1e-6));
484 EXPECT_TRUE(close(std::max(v0, v1), 3.0, 1e-6));
487TEST(LinalgRoutines, GeneralEig) {
488 // Non-symmetric (upper-triangular) [[2,1],[0,3]] -> eigenvalues 2, 3 (real).
489 const la::EigC e = la::eig(mat(2, 2, {2, 1, 0, 3}));
490 EXPECT_TRUE(cclose(cget(e.values, {0}), 3.0)); // descending
491 EXPECT_TRUE(cclose(cget(e.values, {1}), 2.0));
492 // eigvals on the same non-symmetric matrix exercises the general path too.
493 const la::CNDArray ev = la::eigvals(mat(2, 2, {2, 1, 0, 3}));
494 EXPECT_TRUE(cclose(cget(ev, {0}), 3.0));
495 EXPECT_TRUE(cclose(cget(ev, {1}), 2.0));
498TEST(LinalgRoutines, GeneralEigOnSymmetricPromotesToComplex) {
499 // eig() on a symmetric matrix routes through eigh and PROMOTES the real spectrum
500 // and eigenvectors to complex (imag 0): values are real-valued complex, and the
501 // eigenvectors are present (a 2x2 complex matrix), unlike the non-symmetric case.
502 const la::EigC e = la::eig(mat(2, 2, {2, 1, 1, 2})); // eigenvalues 3, 1
503 EXPECT_TRUE(cclose(cget(e.values, {0}), 3.0));
504 EXPECT_TRUE(cclose(cget(e.values, {1}), 1.0));
505 EXPECT_EQ(nd::size_of(e.vectors), 4); // 2x2 eigenvectors present (promoted to complex)
506 EXPECT_EQ(e.vectors.ndim(), 2u);
509// ---- targeted tests for the deep numerical branches ----
511namespace {
512std::vector<double> sorted3(const nd::NDArray& v) {
513 std::vector<double> s{nd::get(v, {0}), nd::get(v, {1}), nd::get(v, {2})};
514 std::sort(s.begin(), s.end());
515 return s;
517// Same, for a complex spectrum known to be real (imaginary parts ≈ 0): the real parts.
518std::vector<double> sorted3c(const la::CNDArray& v) {
519 std::vector<double> s{cget(v, {0}).real(), cget(v, {1}).real(), cget(v, {2}).real()};
520 std::sort(s.begin(), s.end());
521 return s;
523} // namespace
525TEST(LinalgRoutines, VdotInnerAcceptTwoDimVectors) {
526 // A 2-D Nx1 / 1xN is treated as a flat vector by vdot/inner.
527 EXPECT_DOUBLE_EQ(la::vdot(mat(3, 1, {1, 2, 3}), mat(3, 1, {4, 5, 6})), 32.0);
528 EXPECT_DOUBLE_EQ(la::inner(mat(1, 3, {1, 2, 3}), mat(1, 3, {4, 5, 6})), 32.0);
531TEST(LinalgRoutines, DetRequiresPivot) {
532 EXPECT_TRUE(close(la::det(mat(2, 2, {0, 1, 1, 0})), -1.0)); // forces an LU row swap
535TEST(LinalgRoutines, Eigvalsh3x3Dense) {
536 // [[2,1,1],[1,2,1],[1,1,2]] = I + ones -> eigenvalues 4, 1, 1 (Jacobi rotations).
537 const std::vector<double> v = sorted3(la::eigvalsh(mat(3, 3, {2, 1, 1, 1, 2, 1, 1, 1, 2})));
538 EXPECT_TRUE(close(v[0], 1.0, 1e-6));
539 EXPECT_TRUE(close(v[1], 1.0, 1e-6));
540 EXPECT_TRUE(close(v[2], 4.0, 1e-6));
543TEST(LinalgRoutines, GeneralEigvals3x3Dense) {
544 // Same dense matrix through the general (Hessenberg + shifted-QR) path.
545 const std::vector<double> v = sorted3c(la::eigvals(mat(3, 3, {2, 1, 1, 1, 2, 1, 1, 1, 2})));
546 EXPECT_TRUE(close(v[0], 1.0, 1e-6));
547 EXPECT_TRUE(close(v[2], 4.0, 1e-6));
550TEST(LinalgRoutines, NonSymmetric3x3HessenbergPath) {
551 // M = P·diag(2,3,5)·P⁻¹ — non-symmetric with real eigenvalues 2,3,5. Routes
552 // through eigvals_general (Householder–Hessenberg + shifted QR for n≥3).
553 const nd::NDArray M = mat(3, 3, {2.5, 0.5, -0.5, -1, 4, 1, -1.5, 1.5, 3.5});
554 const std::vector<double> v = sorted3c(la::eigvals(M));
555 EXPECT_TRUE(close(v[0], 2.0, 1e-6));
556 EXPECT_TRUE(close(v[1], 3.0, 1e-6));
557 EXPECT_TRUE(close(v[2], 5.0, 1e-6));
558 EXPECT_TRUE(cclose(cget(la::eig(M).values, {0}), 5.0)); // descending; eig() too
561TEST(LinalgRoutines, ComplexEigenvaluesOfRotation) {
562 // A 2-D rotation [[0,-1],[1,0]] has eigenvalues ±i. The general eigensolver
563 // returns the complex conjugate pair (descending by real, then imag: +i, then -i)
564 // rather than throwing — complex spectra are first-class.
565 const la::CNDArray ev = la::eigvals(mat(2, 2, {0, -1, 1, 0}));
566 EXPECT_TRUE(cclose(cget(ev, {0}), std::complex<double>(0.0, 1.0)));
567 EXPECT_TRUE(cclose(cget(ev, {1}), std::complex<double>(0.0, -1.0)));
568 // A complex pair with a non-zero real part: [[1,-1],[1,1]] -> 1±i.
569 const la::CNDArray ev2 = la::eigvals(mat(2, 2, {1, -1, 1, 1}));
570 EXPECT_TRUE(cclose(cget(ev2, {0}), std::complex<double>(1.0, 1.0)));
571 EXPECT_TRUE(cclose(cget(ev2, {1}), std::complex<double>(1.0, -1.0)));
574TEST(LinalgRoutines, VdotRejectsNonVector) {
575 EXPECT_THROW(la::vdot(mat(2, 2, {1, 2, 3, 4}), mat(2, 2, {1, 2, 3, 4})), std::runtime_error);
578TEST(LinalgRoutines, NormOfMatrixIsFrobenius) {
579 EXPECT_TRUE(close(la::norm(mat(2, 2, {1, 2, 2, 4})), 5.0)); // sqrt(1+4+4+16)
582TEST(LinalgRoutines, PinvCondRankOnWideMatrix) {
583 const nd::NDArray W = mat(2, 3, {1, 0, 0, 0, 1, 0}); // 2x3 (more cols than rows)
584 const nd::NDArray P = la::pinv(W); // -> 3x2 (transpose-SVD branch)
585 EXPECT_EQ(nd::shape_of(P), (std::vector<long long>{3, 2}));
586 EXPECT_GE(la::cond(W), 1.0);
587 EXPECT_EQ(la::matrix_rank(W), 2);
590// ---- coverage: non-contiguous inputs, edge branches, and defensive throws ----
591TEST(LinalgRoutines, NonContiguousInputs) {
592 // A broadcast (stride-0) view is non-contiguous, exercising the packing fallback in
593 // contig/as_matrix/as_vector/as_cmatrix (the contiguous fast path runs everywhere else).
594 const nd::NDArray ncvec = nd::broadcast_to(nd::scalar(2.0), {4}); // [2,2,2,2]
595 EXPECT_TRUE(close(la::dot(ncvec, ncvec), 16.0, 1e-12)); // contig() pack path
596 EXPECT_TRUE(close(la::norm(nd::broadcast_to(nd::scalar(3.0), {2, 2})), 6.0, 1e-12));
597 const nd::NDArray ncmat = nd::broadcast_to(nd::array({1.0, 2.0, 3.0}), {3, 3});
598 EXPECT_TRUE(close(la::det(ncmat), 0.0, 1e-9)); // as_matrix pack path
599 const nd::NDArray I3 = mat(3, 3, {1, 0, 0, 0, 1, 0, 0, 0, 1});
600 const nd::NDArray x = la::solve(I3, nd::broadcast_to(nd::scalar(5.0), {3})); // as_vector pack
601 EXPECT_TRUE(close(nd::get(x, {0}), 5.0, 1e-9));
602 const la::CNDArray nccx = nd::broadcast_to(nd::scalar(C(1, 1)), {2, 2});
603 EXPECT_EQ(la::conj_transpose(nccx).ndim(), 2u); // complex contig() pack
604 // complex eigvalsh extracts via as_cmatrix — a broadcast [[2,2],[2,2]] (Hermitian) packs.
605 EXPECT_TRUE(close(nd::get(la::eigvalsh(nd::broadcast_to(nd::scalar(C(2, 0)), {2, 2})), {0}),
606 4.0, 1e-9));
609TEST(LinalgRoutines, ComplexDotFourPlusElements) {
610 // 4+ elements drives the multi-accumulator loop in cdot (both dot and the conjugating vdot).
611 const la::CNDArray a = cvec({C(1, 1), C(2, 0), C(0, 1), C(1, -1), C(2, 2)});
612 const la::CNDArray b = cvec({C(1, 0), C(0, 1), C(1, 1), C(2, 0), C(0, -1)});
613 C dotref{}, vdotref{};
614 for (long long i = 0; i < 5; ++i) {
615 const C ai = cget(a, {i}), bi = cget(b, {i});
616 dotref += ai * bi;
617 vdotref += std::conj(ai) * bi;
618 }
619 EXPECT_TRUE(cclose(la::dot(a, b), dotref)); // non-conjugating multi-accumulator
620 EXPECT_TRUE(cclose(la::vdot(a, b), vdotref)); // conjugating multi-accumulator
623TEST(LinalgRoutines, ShapeAndConvergenceGuards) {
624 const nd::NDArray v = nd::array({1.0, 2.0});
625 EXPECT_THROW(la::matmul(v, v), std::runtime_error); // real matmul non-2D
626 EXPECT_THROW(la::matmul(mat(2, 3, {1, 2, 3, 4, 5, 6}), mat(2, 2, {1, 2, 3, 4})),
627 std::runtime_error); // matmul front inner-dim mismatch (3 != 2)
628 EXPECT_THROW(la::kron(v, v), std::runtime_error); // kron non-2D
629 const la::CNDArray cv = cvec({C(1, 0), C(2, 0)});
630 EXPECT_THROW(la::matmul(cv, cv), std::runtime_error); // complex matmul non-2D
631 // inv that requires a row pivot (zero leading pivot)
632 const nd::NDArray inv = la::inv(mat(2, 2, {0, 1, 1, 0}));
633 EXPECT_TRUE(close(nd::get(inv, {0, 1}), 1.0, 1e-9));
634 // values-only SVD on a WIDE matrix routes through the transpose branch
635 EXPECT_TRUE(close(nd::get(la::svdvals(mat(2, 3, {1, 0, 0, 0, 1, 0})), {0}), 1.0, 1e-9));
636 // NaN input never converges -> the defensive "did not converge" throws fire
637 const double nan = std::numeric_limits<double>::quiet_NaN();
638 EXPECT_THROW(la::eigvalsh(mat(2, 2, {nan, 0, 0, 1})), std::runtime_error); // symmetric QL
639 EXPECT_THROW(la::svdvals(mat(2, 2, {nan, 0, 0, 1})), std::runtime_error); // SVD QR
642TEST(LinalgRoutines, RankDeficientAndDiagonalPaths) {
643 // A matrix with an exactly-zero column gives an exactly-zero singular value, which
644 // drives the g==0 branch in U accumulation and the bulge cancellation in the SVD QR.
645 const la::SVD s = la::svd(mat(3, 3, {1, 4, 0, 2, 5, 0, 3, 6, 0}));
646 EXPECT_TRUE(close(nd::get(s.s, {2}), 0.0, 1e-9));
647 // A diagonal symmetric matrix has all-zero off-diagonals -> the scale==0 branch in tred2.
648 const la::Eig e = la::eigh(mat(3, 3, {2, 0, 0, 0, 5, 0, 0, 0, 7}));
649 EXPECT_TRUE(close(nd::get(e.values, {0}), 7.0, 1e-9));
652TEST(LinalgRoutines, SvdCancellationPath) {
653 // An upper-bidiagonal matrix with an interior zero diagonal (w[1]=0) but a
654 // non-negligible super-diagonal triggers the QR "cancel rv1" Givens sweep in U.
655 (void)la::svd(mat(3, 3, {2, 3, 0, 0, 0, 4, 0, 0, 5}));
656 (void)la::svd(mat(3, 3, {0, 5, 0, 0, 0, 5, 0, 0, 0}));
657 (void)la::svd(mat(4, 4, {1, 9, 0, 0, 0, 0, 9, 0, 0, 0, 0, 9, 0, 0, 0, 1}));
658 SUCCEED();
661// Exercise the WIDE-UNROLL main loops of the multi-accumulator/blocked kernels: the
662// existing tests use tiny matrices that only ever run the scalar remainder, leaving the
663// 4-/8-wide vectorized bodies (ddot, real+complex matmul row-blocking, cholesky/qr/
664// tred2/trace reductions) uncovered. These use n≥8 so the main loops run.
665TEST(LinalgRoutines, WideKernelPaths) {
666 // 8×8 SPD: diag 10, off-diag 1 (= 9·I + J). Eigenvalues {17, 9×7}, trace 80.
667 std::vector<double> a8(64);
668 for (std::size_t i = 0; i < 8; ++i)
669 for (std::size_t j = 0; j < 8; ++j) a8[i * 8 + j] = (i == j) ? 10.0 : 1.0;
670 const nd::NDArray A = mat(8, 8, a8);
672 // dot over ≥8 elements → ddot 8-wide body.
673 EXPECT_DOUBLE_EQ(la::dot(nd::array(std::vector<double>(10, 1.0)),
674 nd::array(std::vector<double>(10, 2.0))), 20.0);
675 // trace 8×8 → trace 4-wide body.
676 EXPECT_DOUBLE_EQ(la::trace(A), 80.0);
677 // matmul 8×8 → real 4-row block. A·I == A.
678 std::vector<double> id8(64, 0.0);
679 for (std::size_t i = 0; i < 8; ++i) id8[i * 8 + i] = 1.0;
680 const nd::NDArray AI = la::matmul(A, mat(8, 8, id8));
681 EXPECT_DOUBLE_EQ(nd::get(AI, {0, 0}), 10.0);
682 EXPECT_DOUBLE_EQ(nd::get(AI, {1, 0}), 1.0);
683 // cholesky 8×8 (j reaches ≥4 → 4-wide inner dot). Reconstruct A = L·Lᵀ.
684 const nd::NDArray L = la::cholesky(A);
685 double a00 = 0;
686 for (long long k = 0; k < 8; ++k) a00 += nd::get(L, {0, k}) * nd::get(L, {0, k});
687 EXPECT_NEAR(a00, 10.0, 1e-9);
688 // qr 8×4 → reflect 4-wide body. R upper-triangular, Q·R == A_panel.
689 std::vector<double> p(32);
690 for (std::size_t i = 0; i < 8; ++i)
691 for (std::size_t j = 0; j < 4; ++j) p[i * 4 + j] = a8[i * 8 + j];
692 const la::QR qr = la::qr(mat(8, 4, p));
693 EXPECT_NEAR(nd::get(qr.r, {1, 0}), 0.0, 1e-9); // upper-triangular
694 // eigvalsh 8×8 → tred2 mat-vec 4-wide body. Largest eigenvalue 17, sum 80.
695 const nd::NDArray w = la::eigvalsh(A);
696 EXPECT_NEAR(nd::get(w, {0}), 17.0, 1e-7);
697 double sw = 0;
698 for (long long i = 0; i < 8; ++i) sw += nd::get(w, {i});
699 EXPECT_NEAR(sw, 80.0, 1e-7);
700 // eig() on a symmetric matrix → the reuse-the-extracted-A symmetric branch.
701 const la::EigC e = la::eig(A);
702 EXPECT_NEAR(cget(e.values, {0}).real(), 17.0, 1e-6);
704 // complex matmul 8×8 → complex 4-row block.
705 std::vector<C> z(64), zi(64, C{0, 0});
706 for (std::size_t i = 0; i < 8; ++i) { z[i * 8 + i] = C{2, 0}; zi[i * 8 + i] = C{1, 0}; }
707 const la::CNDArray Z = cmat(8, 8, z), I = cmat(8, 8, zi);
708 EXPECT_TRUE(cclose(cget(la::matmul(Z, I), {3, 3}), C{2, 0}));
711// Non-contiguous (broadcast/strided) operands take the scratch-packing fallback in the
712// products/reductions, not the zero-copy fast path.
713TEST(LinalgRoutines, StridedOperandFallback) {
714 const nd::NDArray s = nd::broadcast_to(nd::scalar(2.0), {10}); // stride-0 view, len 10
715 EXPECT_DOUBLE_EQ(la::dot(s, s), 40.0); // 10 · (2·2)
716 EXPECT_NEAR(la::norm(s), std::sqrt(40.0), 1e-9);
717 const la::CNDArray cs = nd::broadcast_to(nd::scalar(C{2, 0}), {6});
718 EXPECT_TRUE(cclose(la::dot(cs, cs), C{24, 0})); // 6·(2·2)
719 EXPECT_TRUE(cclose(la::vdot(cs, cs), C{24, 0})); // conj path, strided
722// Batched matmul: [B,M,K] @ [B,K,N] -> [B,M,N], each slice the same product the 2-D kernel
723// gives; strict batching (mismatched batch counts / mixed ranks throw).
724TEST(LinalgRoutines, BatchedMatmul) {
725 // Two batches of 2x3 @ 3x2, second batch = 2x the first: slices must match the 2-D results.
726 std::vector<double> av{1, 2, 3, 4, 5, 6}, bv{7, 8, 9, 10, 11, 12};
727 std::vector<double> abatch(av), bbatch(bv);
728 for (double x : av) abatch.push_back(2 * x); // batch 1 doubles A
729 for (double x : bv) bbatch.push_back(x); // batch 1 reuses B
730 const nd::NDArray A = nd::reshape(nd::array(std::move(abatch)), {2, 2, 3});
731 const nd::NDArray B = nd::reshape(nd::array(std::move(bbatch)), {2, 3, 2});
732 const nd::NDArray C3 = la::matmul(A, B);
733 ASSERT_EQ(C3.ndim(), 3u);
734 EXPECT_EQ(C3.shape()[0], 2u);
735 EXPECT_EQ(C3.shape()[1], 2u);
736 EXPECT_EQ(C3.shape()[2], 2u);
737 // slice 0: the classic {58 64; 139 154}; slice 1 doubles it.
738 EXPECT_DOUBLE_EQ(nd::get(C3, {0, 0, 0}), 58.0);
739 EXPECT_DOUBLE_EQ(nd::get(C3, {0, 1, 1}), 154.0);
740 EXPECT_DOUBLE_EQ(nd::get(C3, {1, 0, 0}), 116.0);
741 EXPECT_DOUBLE_EQ(nd::get(C3, {1, 1, 1}), 308.0);
743 // The out-param kernel form reuses the caller's 3-D buffer.
744 nd::NDArray out = nd::zeros({2, 2, 2});
745 la::matmul(out, A, B);
746 EXPECT_DOUBLE_EQ(nd::get(out, {1, 0, 1}), 2 * 64.0);
748 // Strictness: mixed rank and batch-count mismatch throw.
749 EXPECT_THROW(la::matmul(A, nd::zeros({3, 2})), std::runtime_error);
750 EXPECT_THROW(la::matmul(A, nd::zeros({3, 3, 2})), std::runtime_error);
753// The REAL instantiation of the (conjugate-)transpose kernel. The complex path is covered by
754// ComplexProducts, but a real element takes the other side of the kernel's `if constexpr` — a plain
755// transpose with the conjugation compiled out — and nothing exercised it, so the branch that most
756// users actually hit was the untested one.
757TEST(LinalgRoutines, ConjTransposeOnRealMatrix) {
758 const nd::NDArray M = mat(2, 3, {1, 2, 3, 4, 5, 6});
759 const nd::NDArray T = la::conj_transpose(M);
760 ASSERT_EQ(T.ndim(), 2);
761 EXPECT_EQ(T.shape()[0], 3u);
762 EXPECT_EQ(T.shape()[1], 2u);
763 EXPECT_TRUE(close(nd::get(T, {0, 0}), 1));
764 EXPECT_TRUE(close(nd::get(T, {0, 1}), 4));
765 EXPECT_TRUE(close(nd::get(T, {2, 0}), 3));
766 EXPECT_TRUE(close(nd::get(T, {2, 1}), 6));
768 // Same answer through the allocation-free out-param form.
769 nd::NDArray out = nd::NDArray::uninitialized({3, 2});
770 la::conj_transpose(out, M);
771 EXPECT_TRUE(close(nd::get(out, {0, 1}), 4));
772 EXPECT_TRUE(close(nd::get(out, {2, 0}), 3));
775// The out-param shape guard on the COPY path. `outer` validates its out against an initializer-list
776// shape; the routines that build a result and then copy it in (cholesky, inv, solve, lstsq) validate
777// against a runtime shape vector instead — a separate overload, and one no test reached. A wrong-shaped
778// out must be refused rather than write past the caller's buffer.
779TEST(LinalgRoutines, OutParamRejectsWrongShapeOnTheCopyPath) {
780 const nd::NDArray spd = mat(2, 2, {4, 2, 2, 3}); // symmetric positive-definite
781 nd::NDArray good = nd::NDArray::uninitialized({2, 2});
782 EXPECT_NO_THROW(la::cholesky(good, spd));
784 nd::NDArray wrong = nd::NDArray::uninitialized({3, 3});
785 EXPECT_THROW(la::cholesky(wrong, spd), std::runtime_error);