refactor(trace): compute the CPU Tensor trace in the container layer by IvanaGyro · Pull Request #862 · Cytnx-dev/Cytnx

IvanaGyro · 2026-05-31T03:06:13Z

Part of #834.

Stacked on #860 (stride_view). This PR's diff is only the CPU implementation. #850 builds on this for the GPU side.

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 in an unnamed namespace that derives diagonal_length / output_size / per-surviving-axis layout / the output_is_scalar flag inline from Tn.shape() and Tn.strides(). Trace.cpp is 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; refactor(trace): implement the GPU Tensor trace as CUDA kernels #850 reuses it for the GPU.
Stride-aware diagonal sum: PairwiseSum over the stride_view adaptor (from feat(stride_view): add a strided random-access range adaptor #860) 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. diagonal_length == 0 / output_size == 0 guarded with an early zeroed return.
Tensor::from_storage is used to compose the output: fill a flat result Storage directly, then promote it to a Tensor once the diagonal sums are written, so the storage layer stays at the storage layer until the final return.
Odometer ND walk: walks the output elements in row-major order carrying the diagonal's start offset on an odometer — each step bumps the last surviving-axis index (with carry into earlier axes on wrap) and adjusts the offset by the affected axes' input strides. Avoids the per-element division and modulo of decoding the flat index, and needs no precomputed row-major accumulators.
GPU dispatcher signatures in cuTrace_internal.hpp are kept consistent with the CPU side (the actual GPU implementation lands in refactor(trace): implement the GPU Tensor trace as CUDA kernels #850).

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, 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.

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 below are from openblas-cpu (RelWithDebInfo) at the current branch tip (a9d4866), --benchmark_min_time=0.3s, CPU times in µs (smaller is better). Strided wins at every tested size in every regime — trace touches O(n·middle) diagonal elements while the GEMM baseline touches all O(n²·middle), so BLAS throughput never closes the gap, even in the small-n / large-middle regime where many short dots chain through tuned kernels:

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

n / middle	Strided Double	Matvec Double (GEMM)	Strided ComplexDouble	Matvec ComplexDouble
64 / 64	4.14	665	4.61	840
256 / 64	93.3	22360	146	32352
1024 / 16	233	106312	213	207522
2048 / 16	658	449672	773	827872
4096 / 8	764	927594	787	1692563
64 / 2048	559	41562	715	73370
32 / 8192	1425	48168	1323	75193

The last two rows are the BLAS-favorable shapes (small n, large middle) added at review request: strided still wins by 34–74×.

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

n	Strided Double	Vecdot Double (dot)	Reshape Double	Strided ComplexDouble	Vecdot ComplexDouble	Reshape ComplexDouble
64	0.340	1.68	10.7	0.322	2.53	11.6
256	0.440	5.79	158	0.434	10.0	195
1024	1.06	161	3794	1.12	211	4971
4096	35.7	1668	191056	39.9	5235	306758
8192	97.2	10847	906400	104	29265	1277788

Test plan

openblas-cpu: full suite 1072/1072 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.

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

(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 · 2026-06-09T14:48:52Z

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 · 2026-06-10T11:32:08Z

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

IvanaGyro · 2026-06-16T20:56:53Z

Rebased onto the new mdspan-based #860 (no longer has stride_view); added one follow-up commit 73880c0/9fa5b19 that swaps the trace's strided-view consumer to a Kokkos mdspan. The five existing commits and their inline reviews are preserved.

Diff at the top commit:

pairwise_sum.hpp generalized: PairwiseSumBlocks now drives the reduction through a callable get(i) instead of an iterator's first[i]. The existing PairwiseSum(range) wraps [first](i){ return first[i]; }; a new overload accepts a 1D Kokkos::mdspan and wraps [&mds](i){ return mds(i); }. Both go through the same recursive kernel.
The recursive split passes a base offset through the callable (PairwiseSumBlocks(get, base, n) → PairwiseSumBlocks(get, base + half, n - half)) instead of wrapping the callable in a fresh closure at each depth. Wrapping made every recursion level a new template instantiation; for the floating-point dtype matrix in Sum_internal / Trace_internal, the resulting parameter-pack explosion OOM'd cc1plus. The base-offset variant produces one instantiation per (T, Get) regardless of depth.
TraceImpl constructs the diagonal as a 1D mdspan with layout_stride (extent = diagonal_length, stride = strides[ax1] + strides[ax2]) and calls PairwiseSum(diagonal). The 2D path and the ND-odometer hot loop go through the same call site. stride_view is gone from the includes.

openblas-cpu 1051/1051 on the rebased tip.

Generated by Claude Code

codecov · 2026-06-20T11:15:31Z

Codecov Report

❌ Patch coverage is 52.38095% with 40 lines in your changes missing coverage. Please review.
✅ Project coverage is 30.87%. Comparing base (30359b6) to head (8b2ade2).
⚠️ Report is 3 commits behind head on master.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
src/backend/linalg_internal_cpu/Trace_internal.cpp	49.23%	16 Missing and 17 partials ⚠️
src/linalg/Trace.cpp	20.00%	3 Missing and 1 partial ⚠️
include/utils/checked_cast.hpp	50.00%	1 Missing and 1 partial ⚠️
src/Tensor.cpp	90.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #862      +/-   ##
==========================================
+ Coverage   30.15%   30.87%   +0.71%     
==========================================
  Files         238      229       -9     
  Lines       35429    34757     -672     
  Branches    14745    14409     -336     
==========================================
+ Hits        10685    10730      +45     
+ Misses      17450    16720     -730     
- Partials     7294     7307      +13

Flag	Coverage Δ
cpp	`30.38% <52.38%> (+0.71%)`	⬆️
python	`59.41% <ø> (ø)`

Flags with carried forward coverage won't be shown. Click here to find out more.

Components	Coverage Δ
C++ backend	`32.19% <52.38%> (+0.85%)`	⬆️
Python bindings	`17.28% <ø> (ø)`
Python package	`59.41% <ø> (ø)`

Files with missing lines	Coverage Δ
include/Tensor.hpp	`53.05% <ø> (ø)`
src/Tensor.cpp	`17.69% <90.00%> (+1.30%)`	⬆️
include/utils/checked_cast.hpp	`50.00% <50.00%> (ø)`
src/linalg/Trace.cpp	`25.00% <20.00%> (-11.18%)`	⬇️
src/backend/linalg_internal_cpu/Trace_internal.cpp	`48.48% <49.23%> (+39.05%)`	⬆️

... and 14 files with indirect coverage changes

Continue to review full report in Codecov by Harness.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 30359b6...8b2ade2. Read the comment docs.

🚀 New features to boost your workflow:

📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ac4f5e5d7e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

The CPU Tensor trace built an identity UniTensor and called Contract, so the container-level linalg_internal layer depended on the higher-level UniTensor and Contract APIs -- an inverted layering. Compute the trace directly over storage instead, and tighten the internal Trace API at the same time. * Simplify the dispatch API: every Trace_internal_* / cuTrace_internal_* dispatcher (and the Tracefunc_oii typedef) now takes only (Tn, ax1, ax2) and returns the result Tensor. A small TraceParams helper (backend/linalg_internal_cpu/trace_dispatch.hpp) derives Ndiag, Nelem, accu, remain_rank_id, the reduced shape, and the is_2d flag once per call from shape and axes, so the dispatchers and the _trace_2d / _trace_nd kernels no longer plumb those through individually. Trace.cpp loses its derivation block and is now just validation + a dtype-keyed dispatch call. * Add a public Tensor::strides() getter (declared in include/Tensor.hpp, defined in src/Tensor.cpp): for each logical axis, the storage offset increment per +1 step, derived from the shape and inverse mapper -- the same layout Tensor::at already indexes through. This is what the dispatchers use to sum the diagonal in place. * Replace the body of _trace_2d / _trace_nd with a stride-aware reduction via PairwiseSum over a new range adaptor stride_view (snake_case, with a `r | stride(k)` pipe form; iterator stores the underlying base iterator so it stays valid independent of the view's lifetime). The diagonal stride is strides[ax1] + strides[ax2], so the previous contiguous() copy is dropped, and floating-point traces get the same O(log N * eps) error bound as linalg::Sum. Guard Ndiag == 0 / Nelem == 0 (unsigned (Ndiag - 1) would otherwise underflow into a huge span extent). Drop the UniTensor and Generator includes from the CPU translation unit. * On the GPU side, the cuTrace_internal_* dispatchers in cuTrace_internal.cu are rewritten as new-signature trampolines that derive is_2d (rank == 2) and Ndiag (shape[ax1]) from the input and forward to the existing Contract-based _trace_2d_gpu / _trace_nd_gpu helpers. The helpers themselves keep their bodies and only drop the unused Nelem / accu / remain_rank_id / shape parameters from _trace_nd_gpu's signature. This keeps the GPU TU link-consistent with the new cuTrace_internal.hpp declarations on this commit alone; the native CUDA-kernel rewrite of the GPU helpers lands as a separate commit in the stacked GPU branch. * tests/linalg_test/stride_view_test.cpp exercises the iterator concept model, the pipe form, the iterator-outlives-view case, the non-dividing size-vs-stride branch, the zero-stride throw, and complex strided summation. Tests are reached through src/ added to the test binary's include path, so internal headers can be unit-tested directly. * tests/Tensor_strides_test.cpp checks strides() against Tensor::at across ranks 2-5 permutations and across dtypes. * tests/linalg_test/Trace_test.cpp covers the strided in-place path against the contiguous-clone reference across ranks 2-5, complex and integer dtypes, the rank-2 / rank-N output shape invariants, and swapped axis order. * benchmarks/linalg/Trace_bm.cpp pits the in-place strided Trace against the "collect contiguously and reduce" alternatives Ivana asked for: the 3D contiguous baseline is a BLAS matrix-vector multiplication (tr(A[d,m,d]) over d as {middle, n*n} @ {n*n, 1} against vec(I_n)), and the 2D contiguous baseline is the reshape trick (save A[n-1, n-1], view the rest as {n-1, n+1}, permute+contiguous, Sum the resulting first row, add the saved entry). On openblas-cpu debug, the strided path beats both baselines by orders of magnitude at the tested shapes. * The discovery timeout for gtest_discover_tests is bumped to 120 s so the test binary's ASAN+MKL listing fits on slower runners. Co-authored-by: Claude <noreply@anthropic.com>

The ND CPU trace decoded each flat output index into an input base offset by dividing and taking the modulo against precomputed row-major accumulators on every element. Division and modulo dominate that hot loop, and the accumulator vector only existed to support them. Carry the base offset on an odometer instead: precompute the input stride for each surviving axis once, then for each output element bump the last axis index (carrying into earlier axes on wrap) and adjust the base offset by the affected axes' strides. Advancing is amortized O(1) per element with no division, modulo, or per-element flat-index decode, and the row-major accumulator vector is gone. Also drop the now-unused backend/lapack_wrapper.hpp and utils/utils.hpp includes, left over from the previous Contract-based implementation. Co-authored-by: Claude <noreply@anthropic.com>

TraceImpl<T> allocated the result as a Tensor and then reached through out.storage() to fill it, mixing the container and storage layers. Fill a flat result Storage directly and compose the Tensor with Tensor::from_storage once the diagonal sums are written, so the storage is populated at the storage layer and only promoted to a Tensor (and reshaped to the reduced rank) at the end. Behavior is unchanged: the 2D trace yields a single-element result, the ND trace one element per remaining-rank multi-index, and the Ndiag == 0 / Nelem == 0 guard still returns a zeroed result of the reduced shape. Co-authored-by: Claude <noreply@anthropic.com>

…eImpl Two simplifications to the CPU TraceImpl<T>, neither changing behavior: * Remove the std::min/std::max on the trace axes. linalg::Trace validates upstream that ax1 != ax2 and shape[ax1] == shape[ax2], and both quantities TraceImpl derives from the axes -- the diagonal stride strides[ax1] + strides[ax2] and the set of surviving axes (every axis that is neither ax1 nor ax2) -- are symmetric in the two axes, so their order never mattered. * Build the reduced output shape and the per-output-axis input strides in a single pass over the surviving axes, instead of first recording a remain_rank_id list and then indexing strides[remain_rank_id[x]] in a second loop. The odometer hot loop reads out_strides[x] directly. Rename the locals to snake_case (n_diag, n_elem, out_rank) per the Google C++ style guide. Co-authored-by: Claude <noreply@anthropic.com>

The previous identifiers (n_diag, n_elem, out_rank, out_shape, out_strides, data, out_data, base, index, extent, is_2d) were terse abbreviations that described how each variable was used in the loop rather than what it represents in the trace algorithm. Rename them after the concept each one carries: n_diag -> diagonal_length diag_stride -> diagonal_stride extent -> diagonal_span (storage span first..last + 1) n_elem -> output_size out_rank -> surviving_rank (rank of the non-traced axes) out_shape -> output_shape (target shape for reshape_) out_strides -> surviving_input_stride (input stride per surviving axis) is_2d -> output_is_scalar (rank-2 input -> scalar output) shape_in -> input_shape strides -> input_strides data -> input_data out_data -> output_data base -> diagonal_start_offset (start of this output's diagonal) index -> surviving_index (per-surviving-axis odometer) out_storage -> output_storage i / x -> output_index / axis (loop counters by role) No behavior change. Co-authored-by: Claude <noreply@anthropic.com>

…e=8192} shapes The existing 3D Trace shapes all keep middle small (≤ 64) so the diagonal length n dominates the work; the matvec baseline never has enough small-n GEMM rows to amortize its O(n²·middle) traffic against trace's O(n·middle). Add two shapes with the opposite balance — small n and large middle, where BLAS could plausibly close the gap by chaining many short dot-products through tuned kernels — so the comparison is recorded across both regimes rather than only the strided-favourable one. The pair fits in memory at the existing benchmark scale (vec(I_n) and packed each ≤ 8M elements, ~64 MB for double). Co-Authored-By: Claude <noreply@anthropic.com>

The odometer's wrap step subtracts (output_shape[axis] - 1) * stride from diagonal_start_offset to roll the index back to zero. output_shape was already cytnx_int64 (Tensor::reshape_ takes signed extents), so the old code pulled the operands through an explicit static_cast<cytnx_uint64> to do the subtraction in unsigned space and feed an unsigned offset. Make the surviving-axis strides, the surviving-axis index vector, the diagonal start offset, the output linear index, and the surviving rank all cytnx_int64 instead. The wrap subtraction is now a same-type signed expression with no casts, the inner axis loop steps -1..0 with a natural 'axis >= 0' termination rather than the 'axis-- > 0' unsigned-wrap trick, and the comparison ++index[axis] < output_shape[axis] no longer mixes signednesses. Co-Authored-By: Claude <noreply@anthropic.com>

…dex types Tensor::strides() returned std::vector<cytnx_uint64>, but TraceImpl's odometer walk needs signed offsets (it subtracts during wrap-around). This forced a static_cast at every call site and made it easy to miss one of the index types in the function, since some stayed unsigned and others were upgraded to cytnx_int64 ad hoc. Change Tensor::strides() to return std::vector<cytnx_int64> directly, update TraceImpl to use the signed strides without re-casting, and make diagonal_length, diagonal_stride, and diagonal_span consistently cytnx_int64 instead of mixing in cytnx_uint64. The std::vector<cytnx_int64> surviving_index constructor takes its count by value, so the existing static_cast<std::size_t>(surviving_rank) was redundant once surviving_rank is already a signed type that converts implicitly. Add cytnx::internal::CheckedCastToInt64 in include/utils/checked_cast.hpp for the narrowing conversions that remain: input_shape (and storage sizes in general) are still cytnx_uint64, so converting a shape entry to the signed index type used by TraceImpl can in principle overflow. The helper raises a cytnx_error_msg instead of silently wrapping, and is used both in Tensor::strides() and in TraceImpl's shape-to-index conversions. Tensor_strides_test.cpp's ExpectStridesMatchAt helper is updated to accumulate the offset in cytnx_int64 to match the new return type. Co-Authored-By: Claude <noreply@anthropic.com>

IvanaGyro · 2026-06-30T10:34:49Z

@pcchen May you review this PR again? Or are there someone else can review this PR?

pcchen · 2026-06-30T11:26:07Z

@IvanaGyro yes. I will review again.

pcchen · 2026-06-30T11:33:02Z

Review-thread status check — are all raised issues fixed?

Verified each of the 9 review threads against the actual code on claude/trace-cpu, not just the "resolved" flag. 7 of 9 are genuinely fixed; 1 was declined-by-design; 1 is marked resolved but not actually addressed.

Genuinely fixed in code (7)

Hot-loop / and % → odometer — confirmed: the trace loop carries the start offset on an odometer (++surviving_index[axis] with carry); no division/modulo in the loop.
Leftover includes — confirmed removed (lapack_wrapper.hpp, utils/utils.hpp).
Compose output via Tensor::from_storage — confirmed (return Tensor::from_storage(output_storage);).
Benchmark "large middle, moderate n" case — confirmed added ({64,2048}, {32,8192}) with numbers in the description.
Tracefunc_oii typedef flip / GPU decl mismatch — resolved via the refactor(trace): pass scalar Trace internal args by value and unify naming #861 → refactor(trace): compute the CPU Tensor trace in the container layer #862/refactor(trace): implement the GPU Tensor trace as CUDA kernels #850 chain reshape.
GPU build broken — confirmed: a wrapper in cuTrace_internal.cu restores the GPU branch (see caveat below).

Resolved by declining the change (1) — judgment call, not a code change

"Should throw error instead" — the early-return-zero on diagonal_length == 0 || output_size == 0 (Trace_internal.cpp:58-62) was kept, with the reasoning that linalg::Trace validates upstream and zero-extent is legitimate rather than an error. Defensible, and the reviewer didn't object — but the requested change was intentionally not made.

Marked resolved but NOT actually fixed (1) ⚠️

Zero-extent path has no test. The thread is resolved, but the last comment on it says "Flagging this as still outstanding — happy to add the test if you confirm." Confirmed by inspection: no test in Trace_test.cpp uses a zero-sized dimension (shapes are {4,3,4}, {4,3,4,5}, {3,4,3,2,3}, {6,6}). So the diagonal_length == 0 || output_size == 0 early-out remains completely untested. This one is resolved-in-name-only — recommend reopening it and adding a zero-extent case (or confirming the policy on zero-sized-dim tensors first).

Two caveats

The GPU wrapper reintroduces the layering this PR removes. cuTrace_internal.cu builds a UniTensor::eye and calls Contract, so the low-level GPU trace still depends on UniTensor/Contract (it includes UniTensor.hpp). The Refactor Tensor trace: low-level backend depends on UniTensor (API encapsulation violation) #834 goal of removing that inverted dependency is met only on CPU; the GPU side keeps the inversion as a stopgap until refactor(trace): implement the GPU Tensor trace as CUDA kernels #850 lands the CUDA kernels. (Disclosed in the description.)
PR is Draft, reviewDecision = REVIEW_REQUIRED, mergeStateStatus = BLOCKED — no approving review yet.

Bottom line: the correctness/perf items are properly fixed; the one genuine loose end is the zero-extent test coverage, which is closed without being addressed.

Posted by Claude Code on behalf of @pcchen

pcchen · 2026-06-30T12:03:12Z

@IvanaGyro can you add test for the "Zero-extent path"?

Tensor's own constructor and Tensor::get() (slicing) both reject a 0 in any shape dimension, so the existing test suite never exercised TraceImpl's `diagonal_length == 0 || output_size == 0` early-out branch. Add a ZeroExtentTensor helper that builds a genuinely zero-extent Tensor by composing a zero-element Storage (whose Init does not reject zero length) via Tensor::from_storage and reshape_, then cover the three zero-extent shapes relevant to Trace: - {0,0} traced over (0,1): scalar 0 (empty diagonal sums to 0). - {3,0,0} traced over (1,2): shape {3}, all zero (empty diagonal per surviving index). - {3,0,3} traced over (0,2): shape {0} (zero-length surviving axis produces an empty result). These match NumPy's np.trace convention for the equivalent zero-extent inputs. Co-Authored-By: Claude <noreply@anthropic.com>

IvanaGyro · 2026-06-30T20:23:02Z

@IvanaGyro can you add test for the "Zero-extent path"?

@pcchen Yes, added. And here is the decision about tracing zero-dim tensors: #862 (comment)

pcchen · 2026-07-01T02:27:55Z

The failling check will be fixed in #960. I will approve and merge this.

TraceImplGpu mixed cytnx_uint64 (diagonal_length, diagonal_stride, the surviving-axis shape/stride arrays, surviving_rank, output_size) with the cytnx_int64 Tensor::strides() and output_shape already needed, so values crossed between the two with a static_cast at nearly every assignment. Switches every index/extent in TraceImplGpu, TraceKernel, and DecodeDiagonalStartOffset to cytnx_int64 throughout, mirroring the CPU TraceImpl (Trace_internal.cpp, PR #862): CheckedCastToInt64Gpu is now the only place a uint64 shape value crosses into the signed side, and std::ssize replaces a manual cast for surviving_rank. output_shape also now doubles as the surviving-axis extent array uploaded to the device, since after the cast it holds exactly the same values as the former host_surviving_shape, removing that separate vector. blockIdx/gridDim/threadIdx stay CUDA's native unsigned int, cast once at the point they enter the int64 arithmetic, since they aren't part of the host-side bookkeeping this mirrors and CUDA's execution model defines them as unsigned. Also names the previously-repeated literal `0` stream argument once as a local `stream` variable, used consistently across the cudaMallocAsync/cudaMemcpyAsync/cudaEventRecord/cudaFreeAsync calls, and passes it explicitly to the TraceKernel launch's stream slot (previously implicit) so every operation in this function is demonstrably on the same stream rather than relying on the default stream being the same thing by convention. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

gemini-code-assist Bot reviewed May 31, 2026

View reviewed changes

Comment thread src/backend/linalg_internal_cpu/Trace_internal.cpp Outdated

IvanaGyro mentioned this pull request May 31, 2026

refactor(trace): pass scalar Trace internal args by value and unify naming #861

Closed

3 tasks

IvanaGyro commented May 31, 2026

View reviewed changes

Comment thread src/backend/linalg_internal_interface.hpp

Comment thread tests/linalg_test/Trace_test.cpp

IvanaGyro mentioned this pull request May 31, 2026

refactor(trace): implement the GPU Tensor trace as CUDA kernels #850

Merged

4 tasks

IvanaGyro force-pushed the claude/trace-sig-refactor branch from dd601ed to 92ae99d Compare May 31, 2026 08:37

IvanaGyro force-pushed the claude/trace-cpu branch from 621a730 to 915105b Compare May 31, 2026 08:37

pcchen added this to the v1.1.0 milestone Jun 1, 2026

IvanaGyro force-pushed the claude/trace-sig-refactor branch from 92ae99d to f894e02 Compare June 2, 2026 02:11

IvanaGyro force-pushed the claude/trace-cpu branch 2 times, most recently from cce87ec to 954e366 Compare June 2, 2026 03:00

IvanaGyro changed the base branch from claude/trace-sig-refactor to claude/stride-view June 2, 2026 03:00

IvanaGyro force-pushed the claude/stride-view branch from 1e86067 to 539c057 Compare June 9, 2026 15:20

IvanaGyro force-pushed the claude/trace-cpu branch from 954e366 to 6b5a73a Compare June 9, 2026 15:20

IvanaGyro force-pushed the claude/stride-view branch from 539c057 to 6777742 Compare June 10, 2026 23:25

IvanaGyro force-pushed the claude/trace-cpu branch from f19ceb6 to b7b48b2 Compare June 10, 2026 23:25

IvanaGyro mentioned this pull request Jun 16, 2026

feat(stride_view): add a strided random-access range adaptor #860

Merged

2 tasks

IvanaGyro force-pushed the claude/stride-view branch from 6777742 to 9c8e4c0 Compare June 16, 2026 20:56

IvanaGyro force-pushed the claude/trace-cpu branch from b7b48b2 to 9fa5b19 Compare June 16, 2026 20:56

IvanaGyro force-pushed the claude/stride-view branch from 9c8e4c0 to 87ee88a Compare June 17, 2026 04:12

IvanaGyro force-pushed the claude/trace-cpu branch from 9fa5b19 to 2184f11 Compare June 17, 2026 04:12

IvanaGyro force-pushed the claude/stride-view branch from 87ee88a to ece64c0 Compare June 17, 2026 09:17

IvanaGyro force-pushed the claude/trace-cpu branch 2 times, most recently from dfa0fe9 to 07ee406 Compare June 17, 2026 10:33

Base automatically changed from claude/stride-view to master June 19, 2026 04:57

IvanaGyro force-pushed the claude/trace-cpu branch from 3265c7b to 00d0a8a Compare June 20, 2026 17:41

IvanaGyro force-pushed the claude/trace-cpu branch from 00d0a8a to ac4f5e5 Compare June 20, 2026 18:18

IvanaGyro marked this pull request as ready for review June 20, 2026 18:26

chatgpt-codex-connector Bot reviewed Jun 20, 2026

View reviewed changes

Comment thread src/backend/linalg_internal_gpu/cuTrace_internal.hpp

IvanaGyro and others added 8 commits June 21, 2026 02:29

IvanaGyro force-pushed the claude/trace-cpu branch from ac4f5e5 to 4fe03fc Compare June 21, 2026 02:37

IvanaGyro requested a review from pcchen June 22, 2026 08:46

manuschneider mentioned this pull request Jun 23, 2026

Allow to trace over multiple indicex pairs simultaneously #936

Open

Merge branch 'master' into claude/trace-cpu

f721f2b

pcchen mentioned this pull request Jun 30, 2026

Tracking: Tensor/UniTensor trace refactor (CPU + GPU) — status board #959

Closed

3 tasks

pcchen approved these changes Jul 1, 2026

View reviewed changes

pcchen merged commit 4177e4f into master Jul 1, 2026
18 of 19 checks passed

pcchen deleted the claude/trace-cpu branch July 1, 2026 02:28

Uh oh!

Conversation

IvanaGyro commented May 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Change

Tests

Benchmark

Test plan

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

IvanaGyro left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

IvanaGyro commented Jun 9, 2026

Uh oh!

IvanaGyro commented Jun 10, 2026

Uh oh!

IvanaGyro commented Jun 16, 2026

Uh oh!

codecov Bot commented Jun 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

IvanaGyro commented Jun 30, 2026

Uh oh!

pcchen commented Jun 30, 2026

Uh oh!

pcchen commented Jun 30, 2026

Review-thread status check — are all raised issues fixed?

Genuinely fixed in code (7)

Resolved by declining the change (1) — judgment call, not a code change

Marked resolved but NOT actually fixed (1) ⚠️

Two caveats

Uh oh!

pcchen commented Jun 30, 2026

Uh oh!

IvanaGyro commented Jun 30, 2026

Uh oh!

pcchen commented Jul 1, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

IvanaGyro commented May 31, 2026 •

edited

Loading

codecov Bot commented Jun 20, 2026 •

edited

Loading