[jack-dunham]

Adds qr_compact!, a variant of qr_compact that may overwrite A. It takes the same label, permutation and Val arguments.

When A is a DenseArray and its matricization has at least as many rows as columns, Q is formed in A's memory and only R is allocated. If the bipermutation is not already a reshape of A, the permuted matricization is copied back into A first. Otherwise it falls back to MatrixAlgebraKit.qr_compact! on matricize(A, ...).

This avoids a tensor-sized allocation per factorization when the input is no longer needed, e.g. when factorizing a site tensor in a gate application.

[ITensorBot]

Your PR requires formatting changes to meet the project's style guidelines.
Please run the ITensorFormatter to apply these changes.

Click here to view the suggested changes.

diff --git a/src/factorizations.jl b/src/factorizations.jl
index 758a31e..71af015 100644
--- a/src/factorizations.jl
+++ b/src/factorizations.jl
@@ -184,7 +184,10 @@ function qr_compact!(A, perm_codomain, perm_domain; kwargs...)
         end
         F = MatrixAlgebraKit.qr_compact!(A_mat, (A_mat, similar(A_mat, (n, n))); kwargs...)
     else
-        F = MatrixAlgebraKit.qr_compact!(matricize(A, perm_codomain, perm_domain); kwargs...)
+        F = MatrixAlgebraKit.qr_compact!(
+            matricize(A, perm_codomain, perm_domain);
+            kwargs...
+        )
     end
 
     axes_codomain, axes_domain = bipartition_axes(
diff --git a/test/test_exports.jl b/test/test_exports.jl
index cf4171c..7fa69e5 100644
--- a/test/test_exports.jl
+++ b/test/test_exports.jl
@@ -39,7 +39,8 @@ using Test: @test, @testset
                 :output_axes,
                 :PermutedDims, :permuteddims, :permutedims, :permutedims!,
                 :permutedimsadd!, :permutedimsop, :permutedimsopadd!, :project,
-                :project!, :project_aux, :project_hermitian, :projectto!, :qr_compact, :qr_compact!,
+                :project!, :project_aux, :project_hermitian, :projectto!, :qr_compact,
+                :qr_compact!,
                 :qr_full, :rand_map, :randn_map, :right_null, :right_orth, :right_polar,
                 :scalar, :scale!, :scaleadd!, :ScaledBroadcasted, :select_algorithm,
                 :similar_map,

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.12%. Comparing base (f25bff9) to head (208580d).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #243      +/-   ##
==========================================
+ Coverage   83.93%   84.12%   +0.19%     
==========================================
  Files          31       31              
  Lines         971      989      +18     
==========================================
+ Hits          815      832      +17     
- Misses        156      157       +1     

Flag	Coverage Δ
docs	25.15% <0.00%> (-0.49%)	⬇️

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

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[mtfishman]

Do we need this branch? I.e. do we really care if A is overwritten? I think we can just have a "contract" that A might be overwritten, and you can't rely on it any more, and Q should be obtained from the return value.

[jack-dunham]

We do care that A is strictly overwritten for the following reason:

It allows the intermediate (generated by matricizeopcopy) to be garbage collected. For context, this function is to be used when one really wants to reuse the memory of A to avoid having to manually invoke the garbage collector on it, e.g. in cases where one cannot guarantee cases where all references to A are removed. An example is where A = state[vertex], and one wants to reuse the memory of A immediately.

In more detail: If one must permute A to get it the correct form for QR, then one must allocate a temporary C. If one calls QR on this temporary C directly (writing Q into this array), then there are now two A sized tensors in memory: Q (written into C memory), and the input A. That is fine and unavoidable.

Now, if one does not copy this permuted form back into A before the QR, then this Q returned is additional allocated memory to A now existing outside the scope qr_compact!. In addition, the tensor A exists until it is garbage collected which will not occur when the reference state[vertex] still exists.

If we instead overwrite A (which is allowed by the ! contract) using the copyto!, and QR this, then the intermediate C can get garbage collected when qr_compact! returns, and one is left with only one A sized tensor allocated in the calling scope.

One could remove the if statement and have the body of the if statement apply in all cases, but then one gets an unnecessary allocation and call to copyto! when the indices are already in the correct order.

We can change the contract to might be overwritten, but then this function is not useful to me in the one and only place it's used in the gate application code, and I am not really sure what the advantage of that would be.