refactor(trace): compute the CPU Tensor trace in the container layer

[IvanaGyro]

Part of #834.

Stacked on #849 → stride_view → trace-sig-refactor. This PR's diff is only the CPU implementation. Review/merge the prerequisites first; this PR retargets up the chain as each one merges.

Problem

The low-level Tensor trace built an identity UniTensor and called Contract, so the container-level linalg_internal layer depended on the higher-level UniTensor/Contract APIs — an inverted layering. Compute the trace directly over storage instead, and tighten the internal Trace API at the same time.

Change

• Simplified dispatch API: each Trace_internal_* (CPU) and cuTrace_internal_* (GPU) now takes only (Tn, ax1, ax2) and returns the result Tensor. The per-dtype dispatchers each call a single templated TraceImpl<T> in an unnamed namespace that derives Ndiag / Nelem / accu / remain_rank_id / reduced shape / the 2D-vs-Nd flag inline from Tn.shape() and Tn.strides() — no separate TraceParams struct or dispatcher header. Trace.cpp is now just validation + a dtype-keyed dispatch call.
• Public Tensor::strides() getter (snake_case, mirrors shape() / NumPy's .strides): for each logical axis, the storage stride derived from _shape and _invmapper. The CPU trace calls it; the upcoming GPU PR will share the same getter.
• Stride-aware diagonal sum: PairwiseSum over the stride_view adaptor (from the prerequisite PR) with diagonal stride = strides[ax1] + strides[ax2]. The previous contiguous() copy is dropped; floating-point traces get the same O(log N · ε) error bound as linalg::Sum. Ndiag == 0 / Nelem == 0 guards prevent unsigned underflow in (Ndiag - 1) * diag_stride + 1.
• GPU dispatcher signatures in cuTrace_internal.hpp are kept consistent with the CPU side (the actual GPU implementation lands in the follow-up PR).

Tests

• tests/Tensor_strides_test.cpp — walks every multi-index across rank-2..5 permutations and dtypes, asserting flatten(idx · strides()) == offset at() actually reads (so if strides() ever drifts from the layout at() indexes through, every case fails).
• tests/linalg_test/Trace_test.cpp — strided in-place Trace vs linalg::Trace(t.contiguous(), ...) reference across ranks 2–5, complex+integer dtypes, the rank-2 path (_trace_2d / cuTrace_2d_kernel), output rank invariants (rank-N → rank-(N-2)), and swapped axis order.

Benchmark

benchmarks/linalg/Trace_bm.cpp pits the in-place strided Trace against the "collect contiguously and reduce" baselines:

• 3D: matvec — tr(A)[m] = ⟨vec(I_n), vec(A[:, m, :])⟩ as a {middle, n·n} @ {n·n, 1} BLAS GEMM against vec(I_n).
• 2D: vecdot — tr(A) = ⟨vec(I_n), vec(A)⟩ as a BLAS dot product on the n·n-element flattenings.
• 2D: reshape trick — save A[n-1, n-1], view the rest as {n-1, n+1}, permute + contiguous so the first column becomes a contiguous row, reduce it, add the saved corner. The {n-1, n+1} view reuses the source storage in place (Storage::resize is a no-realloc metadata shrink), so the permute → contiguous gather is the only copy; the size is restored (and the dropped corner written back) afterwards. CPU reduces the contiguous row over the raw buffer (no linalg::Sum/Accessor allocations); GPU wraps the row as a {n-1} tensor and uses linalg::Sum since device memory can't be read host-side.

Every variant runs once before timing and is checked against linalg::Trace; a wrong baseline std::aborts instead of producing fast-but-meaningless numbers.

Numbers (openblas-cpu, single run, µs; smaller is better). Strided wins at every tested size, including the largest — trace touches O(n) diagonal elements while the GEMM/dot baselines touch all O(n²), so BLAS throughput never closes the gap:

3D — {n, middle, n} traced over (0, 2):

n / middle	Strided Double	Matvec Double (GEMM)	Strided ComplexDouble	Matvec ComplexDouble
64 / 64	4.78	833	4.88	1045
256 / 64	96.7	28278	152	44281
1024 / 16	233	134696	239	227744
2048 / 16	660	520290	780	914131
4096 / 8	763	1065279	761	1840246

2D — {n, n} traced over (0, 1):

n	Strided Double	Vecdot Double (dot)	Reshape Double	Strided ComplexDouble	Vecdot ComplexDouble	Reshape ComplexDouble
64	0.410	2.17	12.0	0.392	3.00	13.9
256	0.512	7.17	176	0.513	16.0	282
1024	1.17	119	5529	1.21	204	7324
4096	36.9	1756	242078	42.8	5703	352744
8192	92.8	10966	1041902	110	29162	1494712

Test plan

• openblas-cpu: full suite 963/963 passed.
• mkl-cpu: full suite 963/963 passed.
• Each benchmark variant verified against linalg::Trace before timing (the benchmark abort()s if any baseline disagrees) across all sizes above.
• pre-commit (clang-format-14 + EOF + whitespace) clean.

Draft — opened for review.

---

Generated by Claude Code

[gemini-code-assist]

Code Review

This pull request refactors the linalg::Trace implementation to perform direct, stride-aware diagonal sums instead of falling back to expensive tensor contractions. It introduces a Tensor::strides() method to retrieve the storage strides of a tensor, adds comprehensive unit tests for both strides and trace operations, and includes benchmarks to evaluate performance. A review comment suggests optimizing the multi-index and base offset calculation in the CPU trace hot loop by replacing division and modulo operations with an odometer-like incremental index update.

[IvanaGyro]

I think there should be a case that middle is large and n is normal. In this case Matmul may win.

[IvanaGyro]

Not added yet — flagging it. A "BLAS wins when middle is large and n is moderate" case would be a useful complement to the current shapes (which all have small middle or focus the work in n). Two concerns:

1. The current benchmark already covers n / middle = 64/64, 256/64, 1024/16, 2048/16, 4096/8 — strided wins everywhere, with the BLAS matvec scaling as O(n²·middle) (touches the whole tensor) vs strided O(n·middle). For Matmul to win, you'd need a regime where the BLAS throughput more than offsets the n× overcount; that's something like middle ≫ n.
2. Adding e.g. n=32, middle=8192 would size vec(I_n) = n²·middle ≈ 8M and packed ≈ 8M — feasible, ~64 MB for double, but the benchmark single-process build already builds Sum's pairwise-sum buffers etc. and the openblas-cpu sanitized binary is slow.

If you'd like, I can add n=32 / middle=8192 and n=64 / middle=2048 cases in a follow-up — the question is whether you want the result recorded in the PR description before merging, or it's fine to land as a separate "trace bench: BLAS-favorable shapes" PR. Let me know.

---

Generated by Claude Code

[IvanaGyro]

Should throw error here instead.

[IvanaGyro]

I kept the early-return-zero behavior on Ndiag == 0 || Nelem == 0 rather than turning it into an error, for two reasons:

1. linalg::Trace already validates upstream that shape[ax1] == shape[ax2] and that the axes are in range. The only way to reach the internal dispatcher with Ndiag == 0 is if the common traced size is genuinely zero (a {0, m, 0}-style tensor) — which is a well-defined input, not a programmer error. NumPy treats this the same way: np.trace(np.zeros((0,0))) returns 0, not an exception.
2. The guard exists primarily to prevent (Ndiag - 1) from underflowing cytnx_uint64 into a huge extent value (which would OOB the subsequent std::span read). That's a safety guard, not a policy choice — promoting it to a hard error would conflate "this would crash" with "this input is rejected".

If you'd prefer the strict-rejection policy, the natural place is to add a cytnx_error_msg(shape[ax1] == 0, ...) check in linalg::Trace itself (one level up) — happy to do that in a follow-up if you want trace to reject zero-axis tensors as a matter of API contract. Let me know.

---

Generated by Claude Code

[IvanaGyro]

(Review written by Claude on behalf of @IvanaGyro.)

Nice rewrite — moving the trace into a single stride-aware TraceImpl<T> over the container layer removes the UniTensor::eye + Contract round-trip and the old 2D path's contiguity assumption (rawdata[i*Ldim + i]). I worked through the offset math and it is sound: for any output multi-index, base + (Ndiag-1)*diag_stride lands on a valid element, so the std::span(data+base, extent) end pointer is at most one-past storage — no OOB. Tensor::strides() is the right primitive and Tensor_strides_test validates it against at() across permutations/dtypes. CPU coverage (permuted ranks 2–5, Double/ComplexDouble/Int32, 2D path, output-rank, swapped axes) is solid and all of it passed on both openblas-cuda and mkl-cuda. Two things inline: one stacked-build concern and one test gap. Also a minor note: the Reshape_2D benchmark mutates the shared storage of a const input via resize() and relies on shrink-not-reallocating to restore it — it's verified once against Trace, but it's a fragile trick to leave in a committed benchmark; a comment-backed assertion that capacity was preserved would harden it.

[IvanaGyro]

(Review written by Claude on behalf of @IvanaGyro.)

TraceImpl has a dedicated early-out for Ndiag == 0 || Nelem == 0, but no test drives it, so the zero-extent path (and the shape of its result) is unverified. A regression there — wrong output shape, or a stride computation that divides/indexes into empty storage — would slip through.

Suggest a case with a zero-size traced axis (e.g. {0, 0} -> expect a 1-element result, value 0) and one with a zero-size remaining axis (e.g. {3, 0, 3} traced over 0,2 -> expect shape {0}), asserting both shape and that no read occurs. Cheap, and it locks the Ndiag==0 || Nelem==0 branch.

[IvanaGyro]

Flagging this as still outstanding — happy to add the test if you confirm. The reason I held off: creating a Tensor with a zero-sized dimension via random::random_tensor({0, 0}, ...) is a path with no other test coverage in this codebase, and I wasn't sure whether the policy is "accept zero-axis inputs and return a zero trace" (NumPy-style — what my early-return implements) or "reject them" (which the maintainer's other inline comment on Trace_internal.cpp:39 hints at).

If the policy is accept (the current behavior): I'll add ZeroTracedAxisReturnsZero ({0, 0} → {1} containing 0) and ZeroRemainingAxisReturnsShapedZero ({3, 0, 3} traced over (0,2) → {0}) to LinalgTraceTest.

If the policy is reject: the early-return guard should be lifted into linalg::Trace's upstream validation as a cytnx_error_msg, and the test would be EXPECT_THROW cases instead.

Which way would you like it?

---

Generated by Claude Code

[IvanaGyro]

Pushed the rebased branch (954e366) addressing this round, including the squash you asked for on #861:

• refactor(trace): pass scalar Trace internal args by value and unify naming #861 folded into refactor(trace): compute the CPU Tensor trace in the container layer #862 and refactor(trace): implement the GPU Tensor trace as CUDA kernels #850 per your "merge into" note — claude/trace-cpu now sits directly on claude/stride-view, and the Tracefunc_oii typedef + dispatcher signature changes land here on the first CPU commit. The intermediate sig-refactor PR is closed.
• Odometer hot loop (3ceb1db) — accu/div/mod removed; precomputed out_strides[x] + carry-on-wrap exactly as sketched in the inline thread. Amortized O(1) per output element.
• Tensor::from_storage (954e366) — output composed at the storage layer; both 2D and ND paths return via Tensor::from_storage(out_storage).
• Unused includes (3ceb1db) — backend/lapack_wrapper.hpp and utils/utils.hpp (left from the old Contract-based body) are gone.

Tests on the rebased tip: openblas-cpu 965/965, Trace suite 11/11 (including permuted ranks 2–5, complex+integer dtypes, rank-2 path, output-rank, swapped axes); Tensor::strides() covered by Tensor_strides_test walking every multi-index across permuted ranks 2–5.

Open inline:

• "BLAS wins when middle is large" benchmark case — flagged for follow-up; can add n=32/middle=8192 and n=64/middle=2048 if you want it before merge.
• Ndiag==0/Nelem==0 test — held pending your policy call (NumPy-style return-zero vs strict error). I'll add the test the moment you confirm.
• "Should throw" on zero-axis — replied with the case for keeping return-zero parity with NumPy; reversible if you want the strict-error policy.

---

Generated by Claude Code

[IvanaGyro]

Pushed f19ceb6 — companion rename on the CPU side to keep TraceImpl aligned with the GPU side after the matching review round on #850:

was	now
n_diag	diagonal_length
diag_stride	diagonal_stride
extent	diagonal_span
n_elem	output_size
out_rank	surviving_rank
out_shape, out_strides	output_shape, surviving_input_stride
is_2d	output_is_scalar
shape_in, strides	input_shape, input_strides
data, out_data, out_storage	input_data, output_data, output_storage
base, index	diagonal_start_offset, surviving_index
i, x (loop counters)	output_index, axis

No behavior change. openblas-cpu 968/968 from the rebased tip. The CPU and GPU TUs now share the same vocabulary, which made reviewing the two against each other much easier in this last round.

---

Generated by Claude Code