fix(UniTensor): thread non-mutating block views through contract()

[yingjerkao]

Completes the #724 fix by converting the last remaining shared-block mutation site — the KNOWN-ISSUE remnant documented in-code by the parent PR #998. Stacked on #998 (base = fix/blockut-permute-shared-blocks); the diff here is a single commit.

Problem

BlockUniTensor::contract and BlockFermionicUniTensor::contract permute_()/reshape_()'d the operands' blocks in place before each Matmul/Gemm_Batch call and restored the metadata afterward. Any UniTensor sharing those blocks (e.g. via relabel(), which is documented to share block storage):

• observed transiently corrupted blocks during the contraction,
• was left with replaced, permuted (non-contiguous) block storage afterward even though values were restored, and
• was corrupted permanently if an exception fired before the restore ran.

In particular, contracting a UniTensor with a relabel() of itself threw a rank-mismatch error from permute_ on an already-reshaped shared block.

Fix

In all three branches of both files (UNI_MKL integer-dtype Matmul fallback, UNI_MKL fp/complex Gemm_Batch, non-MKL Matmul fallback), the rank-2 matrix views of the blocks are now built with the non-mutating Tensor::permute()/reshape() into local variables instead of mutating and restoring the blocks. Right-block views are cached per block index (the same right block can pair with several left blocks), preserving the old code's one-copy-per-participating-block cost. The Gemm_Batch eager cast-clone of the right operand (clone_meta/astype/manual delete) is replaced by a lazy astype(common_dtype) when each right block's view is built, which also removes the manually-deleted heap clone from the exception path.

Tests (TDD red/green)

For both BlockUniTensor and BlockFermionicUniTensor:

• ContractAliasedSharedBlocksOperandsIntact — contracts a UniTensor with a relabel() of itself, so the operands alias each other's blocks (including block-with-itself pairs); pre-fix this threw mid-contraction. Asserts the result matches contracting independent clones and that both operands keep their values, shapes, and contiguity.
• ContractLeavesSharedBlockObserverUntouched — a third UniTensor sharing an operand's blocks must observe no change; pre-fix it ended up with replaced, non-contiguous block storage.

No dedicated exception-path test: on non-MKL builds no reachable error fires inside the (now removed) mutation window on valid inputs, and post-fix there is no mutation to leak in the first place.

The UNI_MKL branches are not compiled by the local USE_MKL=OFF build; they were additionally syntax/type-checked with -DUNI_MKL -fsyntax-only.

Full C++ suite on top of #998: 1092 ran / 1077 passed / 11 skipped / 4 known-failing (pre-existing linalg_Test.BkUt_*Svd_truncate_return_err* failures, unrelated to this change).

Merge order: merge this PR into fix/blockut-permute-shared-blocks before #998 merges, or merge #998 first and let GitHub retarget this PR to master. Fixes #724 once both land.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request resolves issue #724 by ensuring that tensor contraction operations treat operand blocks as read-only. Instead of mutating the blocks in place and restoring them afterward, the code now uses non-mutating permute and reshape operations to build matrix views, caching the right-block views to prevent redundant computations. Regression tests have also been added to verify that operands with shared blocks remain intact. The review feedback highlights a potential dtype mismatch issue where Lmat is not cast to common_dtype before being passed to Gemm_Batch, which could lead to runtime errors. Additionally, it is recommended to conditionally cast Rtn->_blocks[b] to common_dtype only when their dtypes differ to avoid unnecessary overhead.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

When this->dtype() != common_dtype (e.g., if this->dtype() is Float and rhs->dtype() is Double), Lmat is not cast to common_dtype. Since Gemm_Batch requires all input matrices to have matching data types, this will lead to a runtime error or undefined behavior due to a dtype mismatch between Lmat and Rmat[b].

We should conditionally cast Lmat to common_dtype if its dtype differs.

              Tensor Lmat = this->_blocks[a].permute(mapperL).reshape({-1, comm_dim});
              if (Lmat.dtype() != common_dtype) {
                Lmat = Lmat.astype(common_dtype);
              }

[gemini-code-assist]

When this->dtype() != common_dtype (e.g., if this->dtype() is Float and rhs->dtype() is Double), Lmat is not cast to common_dtype. Since Gemm_Batch requires all input matrices to have matching data types, this will lead to a runtime error or undefined behavior due to a dtype mismatch between Lmat and Rmat[b].

We should conditionally cast Lmat to common_dtype if its dtype differs.

              Tensor Lmat = this->_blocks[a].permute(mapperL).reshape({-1, comm_dim});
              if (Lmat.dtype() != common_dtype) {
                Lmat = Lmat.astype(common_dtype);
              }

[gemini-code-assist]

Calling astype(common_dtype) unconditionally on Rtn->_blocks[b] can introduce unnecessary overhead (such as copying the tensor data or executing type conversion dispatch) when the block's dtype is already equal to common_dtype.

We should check if the block's dtype differs from common_dtype before calling astype.

Suggested change

	Rmat[b] =
	Rtn->_blocks[b].astype(common_dtype).permute(mapperR).reshape({comm_dim, -1});
	Tensor tmp_b = Rtn->_blocks[b];
	if (tmp_b.dtype() != common_dtype) {
	tmp_b = tmp_b.astype(common_dtype);
	}
	Rmat[b] = tmp_b.permute(mapperR).reshape({comm_dim, -1});

[gemini-code-assist]

Calling astype(common_dtype) unconditionally on Rtn->_blocks[b] can introduce unnecessary overhead (such as copying the tensor data or executing type conversion dispatch) when the block's dtype is already equal to common_dtype.

We should check if the block's dtype differs from common_dtype before calling astype.

Suggested change

	Rmat[b] =
	Rtn->_blocks[b].astype(common_dtype).permute(mapperR).reshape({comm_dim, -1});
	Tensor tmp_b = Rtn->_blocks[b];
	if (tmp_b.dtype() != common_dtype) {
	tmp_b = tmp_b.astype(common_dtype);
	}
	Rmat[b] = tmp_b.permute(mapperR).reshape({comm_dim, -1});

[pcchen]

Code Review

Eliminates the last shared-block mutation site: contract() on BlockUniTensor/BlockFermionicUniTensor used to permute_/reshape_ the operands' blocks in place (then restore), transiently — or permanently, on an exception — corrupting any UniTensor sharing those blocks (e.g. a relabel()). Now it builds rank-2 matrix views via non-mutating permute()/reshape() into locals, across all three branches of both files.

Correctness — verified by reading all 6 branches

• The view transformation is behavior-preserving. permute(mapperL).reshape({-1, comm_dim}) yields the exact tensor value the old permute_/reshape_ produced in place, so Matmul/Gemm_Batch receive identical inputs; the inverse-permute/reshape restore is correctly dropped since originals are never touched. ✅
• Right-block caching is sound. Rmat[b] (guarded by Rmat_ready[b]) is computed once per right block: comm_dim is necessarily equal across all left blocks pairing with a given b (else the contraction is invalid), and it's actually fewer copies than the old restore-per-pairing. ✅
• Lazy astype(common_dtype) per right-block view matches the old eager whole-operand cast, and removes the manually-deleted heap clone — eliminating the exception-path leak the old tmp_Rtn had. ✅
• Subtle spot done right: the Gemm target reshape uses Lmat.shape()[0] / Rmat[b].shape()[1] (the view shapes), not the original block shapes — exactly where a naive port would go wrong. ✅
• Fermionic sign logic is untouched; BlockUniTensor mirrors the fermionic version minus signs. ✅
• Exception safety is now structural — nothing is mutated, so there's no window to leak; removing the dedicated exception test is justified.
• Tests are strong. Both classes get ContractAliasedSharedBlocksOperandsIntact (operands alias via relabel(), incl. block-with-itself pairs; asserts correct result vs independent clones, bit-exact operand intactness at 0.0 tol, and signflip() equality) and ContractLeavesSharedBlockObserverUntouched (a third sharer stays bit-exact + contiguous). These genuinely pin the fix. ✅

Findings

1. 🟠 The two MKL branches per file are runtime-untested. The regression tests and the "1077 passed" local run were USE_MKL=OFF, so they only exercise the non-MKL Matmul branch. The MKL integer-Matmul and MKL Gemm_Batch branches — which received the more intricate edits (cached Rmat[b], lazy astype, the batch loop, the view-shape reshape target) — were only -fsyntax-only checked, and no CI job builds+runs tests with MKL (coverity-scan uses the mkl preset but is a static scan). My manual read gives confidence they're correct, but a logic slip there (stale cached view, wrong reshape dim, dtype) would pass every current gate. Recommend running the new gtests against an MKL build (locally or a CI job) before merge — the Gemm_Batch path especially.

2. 🟡 No BuildAndTest gtest job on this commit. Only wheels ran (compilation verified, tests not). The 4 new regression tests are local-only evidence; a fresh full-test CI run would confirm them.

3. 🟢 Peak-memory note. Rmat retains all participating right-block views for the whole contraction (vs one-at-a-time in the old same-dtype path) — fewer total copies but higher peak. Reasonable trade, and no regression for the dtype-mismatch path (old already cloned all of Rtn); worth a one-line mention.

Verdict

A clean, correct, well-tested completion of the #724 shared-block fix — the transformation is behavior-preserving, the caching and lazy-cast are sound, and the tests pin the aliasing behavior precisely (verified across all six branches by reading). The one real gate is that the MKL branches have no runtime test coverage anywhere; I'd want the new tests run against an MKL build before merge. Approve-with-nits, pending MKL test coverage.

Posted by Claude Code on behalf of @pcchen

[pcchen]

Update to my review — the MKL-coverage nit is already tracked (no new issue needed)

On my earlier finding that the MKL contract branches (integer-Matmul fallback + Gemm_Batch) are runtime-untested (local suite + CI run USE_MKL=OFF): rather than a new issue, this is covered by existing ones —

• Run GitHub action with all allowed configurations and cache the build result #583 ("Run GitHub action with all allowed configurations") already proposes a CPU CI matrix that includes MKL builds. I've added a note there asking that those configs also run ctest (not just build), so the MKL contract paths get runtime coverage.
• Unaddressed review items from merged PR #760 (mixed real/complex BlockUniTensor contract) #855 already tracks MKL-contract-specific test gaps in this exact code (e.g. item 5: "BlockFermionicUniTensor mixed-dtype contract has zero test coverage", plus MKL-path correctness/exception-safety items 1/2/6).

So the nit stands but is not a blocker for this PR: I verified all six branches by reading, the non-MKL branch is exercised by the new regression tests and proves the transformation, and the MKL runtime-coverage gap is tracked in #583/#855. Verdict unchanged — approve-with-nit (tracked).

Posted by Claude Code on behalf of @pcchen

[pcchen]

Approving on code correctness. The non-mutating view transformation is behavior-preserving across all six branches (verified by reading): matrix views via permute()/reshape(), sound right-block caching, lazy astype matching the old eager cast (and killing the exception-path leak), and the Gemm reshape-target correctly taken from the view shapes. Fermionic signs untouched. The two regression tests pin the aliasing fix precisely (operands + observer bit-exact, signflip equality).

Two non-blocking notes before merge:

1. No BuildAndTest on the current checks. This PR's CI last ran while it was behind-base (pre-#1007), so its C++ gtests are local-only evidence. Now that #1007 (the git-identity fix) is merged and this is retargeted to master, a fresh run should actually execute BuildAndTest — worth triggering (e.g. an empty commit / re-run) before merge so the regression tests are CI-verified, not just local.

2. MKL branches lack runtime coverage — tracked (not a new issue): #583 (MKL in the CI matrix; I asked there that it also run ctest) and #855 (MKL-contract test gaps). Read-verified here; the non-MKL branch is exercised by the new tests.

Verdict: approve; recommend a fresh CI run (real BuildAndTest) before landing.

Posted by Claude Code on behalf of @pcchen