Skip to content

Add BufferedBPGateUpdate, BP gate stages with allocator scratch - #11

Open
jack-dunham wants to merge 10 commits into
blocked-message-updatefrom
buffered-bp-gate-update
Open

jack-dunham wants to merge 10 commits into
blocked-message-updatefrom
buffered-bp-gate-update

Conversation

@jack-dunham

Copy link
Copy Markdown
Owner

Stacked on blocked-message-update (#1), which is stacked on bp-gate-stages and jd/quadratic-form.

Summary

  • BufferedBPGateUpdate is an opt-in subalgorithm of BPApplyGate. Its bp_gate_factorize! and bp_gate_restore! absorb the message roots in place, so each vertex holds two tensors of its own size instead of one per message.
  • absorb_matrices!(alg, ψ, matrices...) overwrites ψ with ψ times each operator in turn, alternating between ψ and one scratch tensor of its size. With a TensorOperationsContract, every contraction goes through TensorOperations.
  • The stages use upstream's message_gauge roots (Use the asymmetric Gram root for the BP simple-update gauge ITensor/ITensorNetworksNext.jl#204). Each (x, y) pair is wrapped as operators on the bond name, so absorb_matrices! keeps ψ's index names while the leg holds the rank space.
  • Dense storage only: graded storage, or messages of a wider element type than the tensor, throw an ArgumentError.

The first commit consolidates the buffered parts of several earlier bp-gate-stages commits, rewritten for the current gate API.

Tests

  • BufferedBPGateUpdate matches SimpleBPGateUpdate on a 2×3 grid in Float32, Float64 and ComplexF64, with the default allocator and with a BufferAllocator.
  • apply_operators leaves its input unchanged, and apply_operator! names the new bond as requested.
  • absorb_matrices! matches apply for 0 to 3 matrices, with MatricizeContract and TensorOperationsContract.
  • Graded storage and wider message element types throw.
  • test_aqua.jl, test_apply_operator.jl, test_prod_tensors.jl, test_blockedmessageupdate.jl, test_beliefpropagation.jl, test_normnetwork.jl and test_quadraticformnetwork.jl pass (1266 tests).

🤖 Generated with Claude Code

@github-actions

github-actions Bot commented Oct 8, 2026 •

Copy link
Copy Markdown

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/apply/apply_operators.jl b/src/apply/apply_operators.jl
index dc9ae4b..b3d9697 100644
--- a/src/apply/apply_operators.jl
+++ b/src/apply/apply_operators.jl
@@ -529,7 +529,8 @@ function bp_gate_factorize!(
         subalgorithm::SimpleBPGateUpdate, op::AbstractITensor, state, env, v, w
     )
     roots = [
-        message_gauge(m; alg = subalgorithm.eigh_alg) for m in incoming_messages(env, v => w)
+        message_gauge(m; alg = subalgorithm.eigh_alg) for
+            m in incoming_messages(env, v => w)
     ]
     ψ = absorb_matrices!(subalgorithm.contract_alg, state[v], first.(roots)...)
     bondname = only(intersect(names(state[v]), names(env[w => v])))
diff --git a/test/test_apply_operator.jl b/test/test_apply_operator.jl
index 5f78ff1..6148c68 100644
--- a/test/test_apply_operator.jl
+++ b/test/test_apply_operator.jl
@@ -226,7 +226,8 @@ end
         end
         buffered = BPApplyGate(; subalgorithm = buffered_update)
         simple_state, _ = apply_operator(gates[1], network, complex_env)
-        buffered_state, _ = apply_operator(gates[1], network, complex_env; alg = buffered)
+        buffered_state, _ =
+            apply_operator(gates[1], network, complex_env; alg = buffered)
         @test prod(buffered_state) ≈ prod(simple_state) rtol = rtol
     end
 

@jack-dunham
jack-dunham force-pushed the buffered-bp-gate-update branch from 831dfea to bf5b00d Compare October 8, 2026 18:50
@jack-dunham
jack-dunham changed the base branch from blocked-message-update to jd/quadratic-form October 8, 2026 18:53
@jack-dunham
jack-dunham changed the base branch from jd/quadratic-form to blocked-message-update October 8, 2026 18:53
@jack-dunham
jack-dunham added this pull request to stack #13 October 8, 2026 18:57
@jack-dunham
jack-dunham force-pushed the buffered-bp-gate-update branch from bf5b00d to a19feff Compare October 8, 2026 19:40
@jack-dunham
jack-dunham force-pushed the buffered-bp-gate-update branch 3 times, most recently from 97d1b2e to 4df1de4 Compare October 15, 2026 19:34
@jack-dunham
jack-dunham force-pushed the buffered-bp-gate-update branch from 4df1de4 to 1a4eb44 Compare October 15, 2026 19:36
@jack-dunham
jack-dunham force-pushed the buffered-bp-gate-update branch from 071f16d to da26595 Compare October 19, 2026 14:41
Jack Dunham and others added 10 commits October 20, 2026 14:48
`absorb_matrices!(alg, t, matrices...)` returns `t * matrices[1] * matrices[2] * …` and may
overwrite `t`. The generic method folds `*`. The TensorOperations method, for dense `t` and
matrices with both legs on one axis and `t`'s element type, alternates between `t` and one
scratch tensor taken from `alg`'s allocator and freed on return, keeping each new leg in the
contracted leg's place; anything else takes the generic method.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ct_alg`

`SimpleBPGateUpdate(; contract_alg = nothing)` absorbs the message roots in
`bp_gate_factorize!` and the inverse roots in `bp_gate_restore!` with `absorb_matrices!`.
With a `TensorOperationsContract` on dense storage the absorption reuses the vertex tensor's
memory and one allocator scratch tensor; graded storage, whose roots change the bond's axis,
allocates as before. Tests compare a `TensorOperationsContract`, with the default and a
`BufferAllocator`, against the default contraction, including complex messages and graded
storage.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`Q` then occupies the gauged vertex's buffer, so `bp_gate_restore!` holds that buffer and
the output rather than three factor-sized buffers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The `ITensor` method works out the labels after each absorption and passes arrays
and labels on. Dense arrays take the two-buffer method, which now throws through
`check_input` for a non-square matrix or one of another element type, and writes
outputs of 2^30 elements or more in blocks of the last leg for cuTENSOR 2.7. Other
arrays throw, since TensorOperations reads a graded array element by element.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With the `Diagonal` first, TensorAlgebra allocates a host output for GPU data and
then indexes the device array element by element.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Each operator's vertices are its operator_support in the network, which every
reader can compute when it needs them: step!, and StopWhenVertexRevisited from
the iterate of its state. A network holding only part of the vertices answers
the lookup through has_dimname and dimnamevertices, so apply_operators and
apply_operators! drop their vertices keyword. The check that every operator
input is on a tensor of the network still runs before any operator is applied.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`induced_subgraph_from_vertices` takes the subgraph of the operator's network and keeps the
output→input pairs whose output is on it, since the operator constructor throws for a name on
no vertex.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`BeliefPropagationEnvironmentPreparation(network, env)` and `update_environment` called the
`NormNetwork` constructor, which only accepts an `ITensorNetwork`. `normnetwork` lets a
package return its own norm network type for its network type.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…age gauge

cuSOLVER's divide-and-conquer Hermitian eigensolver (`cusolverDnXsyevd`, MatrixAlgebraKit's
default for a `CuArray`) returns NaN eigenvalues and eigenvectors for some nearly singular
ComplexF32 messages, which `message_gauge` passes into the gate. `eigh_alg = Jacobi()` avoids it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…the vertex's storage

`apply(op, t)` contracts through `MatricizeContract`, which makes a permuted copy of `t` before
the product, so a one-site gate held three copies of the vertex tensor at once. The gate now
calls `mulopadd!` with `contract_alg` (cuTENSOR for `TensorOperationsContract` on a GPU),
reading a copy of `t` and writing the result over `t`.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@jack-dunham
jack-dunham force-pushed the buffered-bp-gate-update branch from 1c7bfb5 to c1db26a Compare October 20, 2026 18:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant