From 0804f3395e85bc320271a6c6152064356b785de9 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 16 Oct 2026 14:58:59 +0800
Subject: [PATCH 1/7] Refactor simple-update helpers and support timestep
 resets

---
 src/algorithms/time_evolution/apply_mpo.jl    | 28 ++++++----
 src/algorithms/time_evolution/get_cluster.jl  | 51 ++++++++++---------
 src/algorithms/time_evolution/simpleupdate.jl | 32 +++++++-----
 .../time_evolution/simpleupdate3site.jl       | 16 +++---
 test/testsuite/timeevol/timestep.jl           | 17 +++++++
 5 files changed, 89 insertions(+), 55 deletions(-)

diff --git a/src/algorithms/time_evolution/apply_mpo.jl b/src/algorithms/time_evolution/apply_mpo.jl
index c4a61b68d..511efb47d 100644
--- a/src/algorithms/time_evolution/apply_mpo.jl
+++ b/src/algorithms/time_evolution/apply_mpo.jl
@@ -247,15 +247,28 @@ function _flip_virtuals!(
         Ms::Vector{T}, flips::Vector{Bool}; inv::Bool = false
     ) where {T <: GenericMPSTensor}
     @assert length(flips) == length(Ms) - 1
-    for (n, flip) in enumerate(flips)
-        !flip && continue
+    for (n, fl) in enumerate(flips)
+        !fl && continue
         M1, M2 = Ms[n], Ms[n + 1]
-        Ms[n] = TensorKit.flip(M1, numind(M1); inv)
-        Ms[n + 1] = TensorKit.flip(M2, 1; inv)
+        Ms[n] = flip(M1, numind(M1); inv)
+        Ms[n + 1] = flip(M2, 1; inv)
     end
     return Ms
 end
 
+"Apply a bond projector to the left leg of an MPS tensor."
+function _apply_proj_left(Pb::MPSBondTensor, m::GenericMPSTensor{S, N}) where {S, N}
+    pP = ((1,), (2,))
+    pM = ((1,), ntuple(i -> i + 1, N))
+    pPM = (codomainind(m), domainind(m))
+    return tensorcontract(Pb, pP, false, m, pM, false, pPM)
+end
+
+"Apply a bond projector to the right leg of an MPS tensor."
+function _apply_proj_right(m::GenericMPSTensor, Pa::MPSBondTensor)
+    return m * twistdual(Pa, 1)
+end
+
 """
 Find projectors to truncate internal bonds of the cluster `Ms`.
 """
@@ -266,11 +279,8 @@ function _cluster_truncate!(
     # apply projectors
     # M1 -- (Pa1,wt1,Pb1) -- M2 -- (Pa2,wt2,Pb2) -- M3
     for (i, (Pa, Pb)) in enumerate(zip(Pas, Pbs))
-        Ms[i] = Ms[i] * twistdual(Pa, 1)
-        pP = ((1,), (2,))
-        pM = ((1,), ntuple(i -> i + 1, numind(eltype(Ms)) - 1))
-        pPM = (codomainind(Ms[i + 1]), domainind(Ms[i + 1]))
-        Ms[i + 1] = tensorcontract(Pb, pP, false, Ms[i + 1], pM, false, pPM)
+        Ms[i] = _apply_proj_right(Ms[i], Pa)
+        Ms[i + 1] = _apply_proj_left(Pb, Ms[i + 1])
     end
     return wts, ϵs, Pas, Pbs
 end
diff --git a/src/algorithms/time_evolution/get_cluster.jl b/src/algorithms/time_evolution/get_cluster.jl
index 452c0cd3b..c9a97c629 100644
--- a/src/algorithms/time_evolution/get_cluster.jl
+++ b/src/algorithms/time_evolution/get_cluster.jl
@@ -50,40 +50,41 @@ function _nn_vec_direction(nn_vec::CartesianIndex{2})
 end
 
 """
-Given `site1`, `site2` connected by a nearest neighbor bond,
-return the bond index and whether it is reversed from the
-standard orientation (`site1` on the west/south of `site2`).
+Given `site1`, `site2` connected by a nearest neighbor bond, return `(dir, r, c)` where `dir` ∈ {NORTH, EAST, SOUTH, WEST} is the direction from `site1` to `site2`.
+`(r, c)` is the position of the west site for x-bonds (`dir == EAST` or `dir == WEST`) and the south site for y-bonds (`dir == NORTH` or `dir == SOUTH`).
 """
-function _nn_bondrev(site1::CartesianIndex{2}, site2::CartesianIndex{2})
-    diff = site1 - site2
-    if diff == CartesianIndex(0, -1)
+function _nn_bonddir(site1::CartesianIndex{2}, site2::CartesianIndex{2})
+    diff = site2 - site1
+    if diff == CartesianIndex(0, 1)
         r, c = site1[1], site1[2]
-        return (1, r, c), false
-    elseif diff == CartesianIndex(0, 1)
-        r, c = site2[1], site2[2]
-        return (1, r, c), true
-    elseif diff == CartesianIndex(1, 0)
-        r, c = site1[1], site1[2]
-        return (2, r, c), false
+        return (EAST, r, c)
     elseif diff == CartesianIndex(-1, 0)
+        r, c = site1[1], site1[2]
+        return (NORTH, r, c)
+    elseif diff == CartesianIndex(1, 0)
+        r, c = site2[1], site2[2]
+        return (SOUTH, r, c)
+    elseif diff == CartesianIndex(0, -1)
         r, c = site2[1], site2[2]
-        return (2, r, c), true
+        return (WEST, r, c)
     else
         error("`site1` and `site2` are not nearest neighbors.")
     end
 end
 
-function _bond_rotation(x, bonddir::Int, rev::Bool; inv::Bool = false)
-    return if bonddir == 1 # x-bond
-        rev ? rot180(x) : x
-    elseif bonddir == 2 # y-bond
-        if rev
-            inv ? rotr90(x) : rotl90(x)
-        else
-            inv ? rotl90(x) : rotr90(x)
-        end
-    else
-        error("`bonddir` must be 1 (for x-bonds) or 2 (for y-bonds).")
+"""
+Apply the tensor rotation that maps a bond from orientation `dir_from` to orientation `dir_to`, where the directions are one of `NORTH`, `EAST`, `SOUTH`, `WEST`.
+"""
+function _bond_rotation(x, dir_from::Int, dir_to::Int)
+    delta = mod(dir_to - dir_from, 4)
+    return if delta == 0
+        x
+    elseif delta == 1
+        rotr90(x)
+    elseif delta == 2
+        rot180(x)
+    else # delta == 3
+        rotl90(x)
     end
 end
 
diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 5133f428a..fb64c0f47 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -86,9 +86,8 @@ function _su_iter!(
     Ms, open_vaxs, = _get_cluster(state, sites)
     _absorb_weight!(Ms, sites, open_vaxs, env)
     # rotate
-    bond, rev = _nn_bondrev(sites...)
-    dir = first(bond)
-    A, B = _bond_rotation.(Ms, dir, rev; inv = false)
+    (dir, r, c) = _nn_bonddir(sites...)
+    A, B = _bond_rotation.(Ms, dir, EAST)
     # apply gate
     ϵ = 0.0
     local s
@@ -102,9 +101,9 @@ function _su_iter!(
         alg.purified && break # only apply gate to 1st physical leg
     end
     # rotate back
-    A = _bond_rotation(A, dir, rev; inv = true)
-    B = _bond_rotation(B, dir, rev; inv = true)
-    rev && (s = transpose(s))
+    A = _bond_rotation(A, EAST, dir)
+    B = _bond_rotation(B, EAST, dir)
+    (dir in (WEST, SOUTH)) && (s = transpose(s))
     # remove environment weights
     siteA, siteB = sites
     A = absorb_weight(A, env, siteA[1], siteA[2], open_vaxs[1]; inv = true)
@@ -112,7 +111,8 @@ function _su_iter!(
     # update tensor dict and weight on current bond
     state[siteA] = normalize!(A, Inf)
     state[siteB] = normalize!(B, Inf)
-    env[bond...] = normalize!(s, Inf)
+    d = dir in (EAST, WEST) ? 1 : 2
+    env[d, r, c] = normalize!(s, Inf)
     return ϵ
 end
 
@@ -128,10 +128,11 @@ function su_iter(
         if length(sites) == 1
             # 1-site gate
             # TODO: special treatment for bipartite state
-            site = sites[1]
+            site = only(sites)
             state2[site] = _apply_sitegate(state2[site], gate; alg.purified)
         elseif length(sites) == 2
-            (d, r, c), = _nn_bondrev(sites...)
+            (dir, r, c) = _nn_bonddir(sites...)
+            d = dir in (EAST, WEST) ? 1 : 2
             alg.bipartite && iseven(r) && continue
             ϵ′ = _su_iter!(state2, gate, env2, sites, alg)
             ϵ = max(ϵ, ϵ′)
@@ -169,17 +170,20 @@ end
 
 """
     timestep(
-        it::TimeEvolver{<:SimpleUpdate}, psi::InfiniteState, env::SUWeight
+        it::TimeEvolver{<:SimpleUpdate}, psi::InfiniteState, env::SUWeight;
+        iter::Int = it.state.iter, t::Number = it.state.t
     ) -> (psi, env, info)
 
-Given the `TimeEvolver` iterator `it`, perform one step of time evolution
-on the input state `psi` and its environment `env`.
+Given the `TimeEvolver` iterator `it`, perform one step of time evolution on the input state `psi` and its environment `env`.
+
+Use `iter` and `t` to reset the current iteration number and evolved time of `it`, respectively.
 """
 function MPSKit.timestep(
-        it::TimeEvolver{<:SimpleUpdate}, psi::InfiniteState, env::SUWeight
+        it::TimeEvolver{<:SimpleUpdate}, psi::InfiniteState, env::SUWeight;
+        iter::Int = it.state.iter, t::Number = it.state.t
     )
     _timeevol_sanity_check(psi, physicalspace(it.state.psi), it.alg)
-    state = SUState(it.state.iter, it.state.t, psi, env)
+    state = SUState(iter, t, psi, env)
     result = iterate(it, state)
     if result === nothing
         @warn "TimeEvolver `it` has already reached the end."
diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl
index 2b5b87673..e29d91e8f 100644
--- a/src/algorithms/time_evolution/simpleupdate3site.jl
+++ b/src/algorithms/time_evolution/simpleupdate3site.jl
@@ -47,13 +47,14 @@ function _su_iter!(
     # restore virtual arrows in `Ms`
     _flip_virtuals!(Ms, flips)
     # update env weights
-    bond_revs = map(sites, Iterators.drop(sites, 1)) do site1, site2
-        _nn_bondrev(site1, site2)
+    bond_dirs = map(sites, Iterators.drop(sites, 1)) do site1, site2
+        _nn_bonddir(site1, site2)
     end
-    for (wt, (bond, rev), flip) in zip(wts, bond_revs, flips)
+    for (wt, (dir, r, c), flip) in zip(wts, bond_dirs, flips)
         wt_new = flip ? _fliptwist_s!(wt) : wt
-        wt_new = rev ? transpose(wt_new) : wt_new
-        env[CartesianIndex(bond)] = normalize!(wt_new, Inf)
+        wt_new = (dir in (WEST, SOUTH)) ? transpose(wt_new) : wt_new
+        d = dir in (EAST, WEST) ? 1 : 2
+        env[d, r, c] = normalize!(wt_new, Inf)
     end
     # update state tensors
     for (M, s, invperm, vaxs) in zip(Ms, sites, invperms, open_vaxs)
@@ -72,9 +73,10 @@ function _get_cluster_trunc(
         trunc::TruncationStrategy, sites::Vector{CartesianIndex{2}}
     )
     return map(sites, Iterators.drop(sites, 1)) do site1, site2
-        (d, r, c), rev = _nn_bondrev(site1, site2)
+        (dir, r, c) = _nn_bonddir(site1, site2)
+        d = dir in (EAST, WEST) ? 1 : 2
         t = truncation_strategy(trunc, d, r, c)
-        if rev && isa(t, TruncationSpace)
+        if (dir in (WEST, SOUTH)) && isa(t, TruncationSpace)
             t = truncspace(flip(t.space)')
         end
         return t
diff --git a/test/testsuite/timeevol/timestep.jl b/test/testsuite/timeevol/timestep.jl
index 210a7185f..025c3f981 100644
--- a/test/testsuite/timeevol/timestep.jl
+++ b/test/testsuite/timeevol/timestep.jl
@@ -32,5 +32,22 @@ function timeevol_timestep(AT)
         @test ψ1 == ψ2 == ψ3
         @test env1 == env2 == env3
         @test info1 == info2 == info3
+
+        # Restart an exhausted evolver using the public timestep keywords.
+        ψreset, envreset, inforeset = timestep(evolver, ψ0, env0; iter = 0, t = 0.0)
+        fresh = TimeEvolver(ψ0, H, dt, nstep, alg, env0)
+        ψfresh, envfresh, infofresh = timestep(fresh, ψ0, env0)
+        @test ψreset == ψfresh
+        @test envreset == envfresh
+        @test inforeset == infofresh
+        @test evolver.state.iter == 1
+        @test evolver.state.t == dt
+
+        # Iteration and time offsets affect bookkeeping, not the update.
+        ψoffset, envoffset, infooffset = timestep(evolver, ψ0, env0; iter = 7, t = 1.25)
+        @test ψoffset == ψfresh
+        @test envoffset == envfresh
+        @test evolver.state.iter == 8
+        @test evolver.state.t == infooffset.t == 1.25 + dt
     end
 end

From 7fe751d4cedc214661fa1f472168afca40cd64c4 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 16 Oct 2026 15:03:05 +0800
Subject: [PATCH 2/7] Keep timestep reset coverage within the existing
 equivalence test

---
 test/testsuite/timeevol/timestep.jl | 22 ++--------------------
 1 file changed, 2 insertions(+), 20 deletions(-)

diff --git a/test/testsuite/timeevol/timestep.jl b/test/testsuite/timeevol/timestep.jl
index 025c3f981..7a125b435 100644
--- a/test/testsuite/timeevol/timestep.jl
+++ b/test/testsuite/timeevol/timestep.jl
@@ -22,9 +22,8 @@ function timeevol_timestep(AT)
         # time_evolve
         ψ2, env2, info2 = time_evolve(ψ0, H, dt, nstep, alg, env0)
         # for-loop syntax
-        ## manually reset internal state of evolver
-        evolver.state = PEPSKit.SUState(0, 0.0, ψ0, env0)
-        ψ3, env3, info3 = nothing, nothing, nothing
+        # Restart the exhausted evolver using the public timestep keywords.
+        ψ3, env3, info3 = timestep(evolver, ψ0, env0; iter = 0, t = 0.0)
         for state in evolver
             ψ3, env3, info3 = state
         end
@@ -32,22 +31,5 @@ function timeevol_timestep(AT)
         @test ψ1 == ψ2 == ψ3
         @test env1 == env2 == env3
         @test info1 == info2 == info3
-
-        # Restart an exhausted evolver using the public timestep keywords.
-        ψreset, envreset, inforeset = timestep(evolver, ψ0, env0; iter = 0, t = 0.0)
-        fresh = TimeEvolver(ψ0, H, dt, nstep, alg, env0)
-        ψfresh, envfresh, infofresh = timestep(fresh, ψ0, env0)
-        @test ψreset == ψfresh
-        @test envreset == envfresh
-        @test inforeset == infofresh
-        @test evolver.state.iter == 1
-        @test evolver.state.t == dt
-
-        # Iteration and time offsets affect bookkeeping, not the update.
-        ψoffset, envoffset, infooffset = timestep(evolver, ψ0, env0; iter = 7, t = 1.25)
-        @test ψoffset == ψfresh
-        @test envoffset == envfresh
-        @test evolver.state.iter == 8
-        @test evolver.state.t == infooffset.t == 1.25 + dt
     end
 end

From 194954946dddd09cdaee73795afa73afa73164b8 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 16 Oct 2026 16:25:34 +0800
Subject: [PATCH 3/7] Make `absorb_weight` more flexible in tensor leg order

---
 src/algorithms/contractions/absorb_weight.jl  | 43 ++++++++++++-------
 .../time_evolution/simpleupdate3site.jl       | 16 ++++---
 src/utility/util.jl                           |  4 +-
 3 files changed, 40 insertions(+), 23 deletions(-)

diff --git a/src/algorithms/contractions/absorb_weight.jl b/src/algorithms/contractions/absorb_weight.jl
index f6b764a5a..6362b1aca 100644
--- a/src/algorithms/contractions/absorb_weight.jl
+++ b/src/algorithms/contractions/absorb_weight.jl
@@ -48,18 +48,41 @@ function absorb_weight(
         t::Union{PEPSTensor, PEPOTensor}, weights::SUWeight,
         row::Int, col::Int, virt_axes::NTuple{N, Int}; inv::Bool = false
     ) where {N}
-    Np = numout(t)
+    return _absorb_weight(t, weights, row, col, virt_axes; inv)
+end
+
+"""
+Absorb weights on the virtual directions in `virt_axes`, or remove them when `inv = true`.
+
+- `virt_axes` selects virtual directions using `1 = NORTH`, `2 = EAST`, etc., regardless of the current leg order.
+- `input_axes` specifies how the legs of the input `t` are already permuted with respect to the standard order.
+  For example, `input_axes = (5, 1, 2, 4, 3)` for a PEPS tensor means the legs are currently ordered as west, physical, north, south, east.
+- `output_perm` specifies how the returned tensor should permute the legs with respect to the standard order.
+  For example, `output_perm = ((5, 1, 2, 4), (3,))` puts west, physical, north, south in the codomain and east in the domain.
+- There are always four virtual legs on the square lattice, so the physical leg count is `numind(t) - 4`.
+"""
+function _absorb_weight(
+        t::AbstractTensorMap, weights::SUWeight, row::Int, col::Int,
+        virt_axes::NTuple{N, Int}; inv::Bool = false,
+        input_axes::Tuple = ntuple(identity, numind(t)),
+        output_perm::Tuple{Tuple, Tuple} = (
+            ntuple(identity, numind(t) - 4), ntuple(i -> numind(t) - 4 + i, 4),
+        )
+    ) where {N}
+    Np = numind(t) - 4
     vax = first(virt_axes)
+    axes, biperm = _permute_to_last(input_axes, vax + Np)
     wt = weight_to_absorb(weights, row, col, vax; inv)
-    axes, t2 = absorb_first_weight(t, wt, vax)
+    # Allocate the first result to preserve the input and allow scalar promotion.
+    t2 = permute(t, biperm) * wt
     for vax in Base.tail(virt_axes)
         axes, biperm = _permute_to_last(axes, vax + Np)
         wt = weight_to_absorb(weights, row, col, vax; inv)
-        # use `*` to make absorption/removal twist-free
-        t2 = permute(t2, biperm) * wt
+        # Diagonal composition is twist-free and only mutates the intermediate.
+        t2 = rmul!(permute(t2, biperm), wt)
     end
     perm_back = invperm(axes)
-    return permute(t2, (perm_back[begin:numout(t)], perm_back[(numout(t) + 1):end]))
+    return permute(t2, map(p -> map(i -> perm_back[i], p), output_perm))
 end
 
 """
@@ -86,13 +109,3 @@ function weight_to_absorb(
     (ax == SOUTH || ax == WEST) && return transpose(wt)
     return wt
 end
-
-function absorb_first_weight(
-        t::Union{PEPSTensor, PEPOTensor}, wt::DiagonalTensorMap, vax::Int
-    )
-    Np = numout(t)
-    axes, biperm = _permute_to_last(ntuple(identity, numind(t)), vax + Np)
-    # use `*` to make absorption/removal twist-free
-    t2 = permute(t, biperm) * wt
-    return axes, t2
-end
diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl
index e29d91e8f..a0f55c693 100644
--- a/src/algorithms/time_evolution/simpleupdate3site.jl
+++ b/src/algorithms/time_evolution/simpleupdate3site.jl
@@ -13,14 +13,19 @@ function _unfuse_physicalspace(
     return O_unfused, F
 end
 
+"""
+Absorb open-leg weights directly into the cluster's MPS layout.
+"""
 function _get_cluster_with_weights(
         state::InfiniteState, sites::Vector{CartesianIndex{2}}, env::SUWeight
     )
     Ms, open_vaxs, perms = _get_cluster(state, sites)
-    _absorb_weight!(Ms, sites, open_vaxs, env; inv = false)
     Np = (state isa InfinitePEPS) ? Val(1) : Val(2)
+    Ms_mps = map(Ms, sites, open_vaxs, perms) do M, site, vaxs, perm
+        return _absorb_weight(M, env, site[1], site[2], vaxs; output_perm = perm)
+    end
     invperms = map(p -> _inv_mpo_perm(p, Np), perms)
-    return _permute_cluster(Ms, perms), open_vaxs, invperms
+    return Ms_mps, open_vaxs, invperms
 end
 
 """
@@ -57,9 +62,10 @@ function _su_iter!(
         env[d, r, c] = normalize!(wt_new, Inf)
     end
     # update state tensors
-    for (M, s, invperm, vaxs) in zip(Ms, sites, invperms, open_vaxs)
-        # restore original axes order and remove weights on open axes of the cluster
-        M′ = absorb_weight(permute(M, invperm), env, s, vaxs; inv = true)
+    for (M, s, perm_back, vaxs) in zip(Ms, sites, invperms, open_vaxs)
+        # Remove weights directly from MPS order, then restore the site's layout.
+        input_axes = invperm((perm_back[1]..., perm_back[2]...))
+        M′ = _absorb_weight(M, env, s[1], s[2], vaxs; inv = true, input_axes)
         state[s] = normalize!(M′, Inf)
     end
     return ϵ
diff --git a/src/utility/util.jl b/src/utility/util.jl
index 4c450ddb4..9893262d3 100644
--- a/src/utility/util.jl
+++ b/src/utility/util.jl
@@ -169,8 +169,6 @@ function random_dual!(Vs::AbstractMatrix{E}; p = 0.7) where {E <: ElementarySpac
 end
 
 """
-    _permute_to_last(axes::NTuple{N, Int}, ax::Int) where {N}
-
 Returns `(1, 2, ..., N)` but with `ax` moved to the end,
 and the corresponding permutation for `axes` (with `ax` as the only domain index).
 """
@@ -178,6 +176,6 @@ function _permute_to_last(axes::NTuple{N, Int}, ax::Int) where {N}
     codomain_axes = TupleTools.deleteat(ntuple(identity, N), ax)
     q = invperm(axes)
     biperm = (map(i -> q[i], codomain_axes), (q[ax],))
-    new_axes = (ntuple(i -> axes[biperm[1][i]], N - 1)..., ax)
+    new_axes = (codomain_axes..., ax)
     return new_axes, biperm
 end

From 7be1ad8b103731fb1c5cb5454afa0d0c7960ee6d Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Wed, 28 Oct 2026 18:12:48 -0400
Subject: [PATCH 4/7] Rotation-free simple update kernel with buffered
 temporaries

Rework the nearest-neighbor simple update along the lines of the Canopy.jl
two-site kernel:

- Absorb diagonal bond weights directly on each fusion-tree block, in the
  native leg order, with a single fused Strided broadcast. This replaces the
  permute-multiply-permute loop and the `input_axes`/`output_perm` keyword
  arguments of `_absorb_weight`.
- Update a bond without rotating the tensors: derive the legs from the bond
  direction, split each site with one permutation and a QR (`_bond_qr`), apply
  the gate to the reduced tensors (`_apply_bond_gate`), and contract the result
  straight back into the native leg order (`_undo_bond_qr`). Orient the bond
  such that the SVD reproduces the original arrow, so no flips are needed.
- Serve temporaries from `MPSKit.default_allocator`, created once per
  `su_iter`. Temporaries never outlive the function that allocates them.

Remove the now unused `_bond_rotation`, `_absorb_weight!`, `_permute_cluster`,
`_permute_to_last` and `_apply_gate`, add `absorb_weight!` and `absorb_s!`,
and require TensorOperations 5.8.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/PEPSKit.jl                                |   5 +-
 src/algorithms/contractions/absorb_weight.jl  |  67 ++++++----
 src/algorithms/time_evolution/apply_gate.jl   | 123 ++++++++++++++----
 src/algorithms/time_evolution/get_cluster.jl  |  40 ------
 src/algorithms/time_evolution/simpleupdate.jl |  57 ++++----
 .../time_evolution/simpleupdate3site.jl       |  18 ++-
 src/utility/util.jl                           |  18 +--
 7 files changed, 192 insertions(+), 136 deletions(-)

diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl
index 2aeb865b6..e3b683987 100644
--- a/src/PEPSKit.jl
+++ b/src/PEPSKit.jl
@@ -23,6 +23,8 @@ using KrylovKit
 using KrylovKit: Lanczos, BlockLanczos
 
 using TensorOperations, OptimKit
+using TensorOperations: DefaultAllocator, DefaultBackend, tensoralloc_add
+using TensorKit.Strided: StridedView, sreshape
 import OptimKit: fixedpoint
 using ChainRulesCore, Zygote
 import TupleTools
@@ -32,6 +34,7 @@ using MPSKit: MPSTensor, MPOTensor, GenericMPSTensor, MPSBondTensor, ProductTran
 using MPSKit: InfiniteEnvironments
 using MPSKit: DynamicTol
 using MPSKit.DynamicTols: adapt_solver
+using MPSKit: default_allocator
 using MPSKit: @log_initialization, @log_iteration, @log_convergence, @log_nonconvergence
 using MPSKit.IterativeLoggers: with_verbosity
 import MPSKit.DynamicTols: _updatetol
@@ -49,7 +52,7 @@ import MPSKitModels:
     transverse_field_ising, heisenberg_XYZ, heisenberg_XXZ,
     hubbard_model, bose_hubbard_model, tj_model
 using FiniteDifferences
-using OhMyThreads: tmap, tmap!
+using OhMyThreads: tmap, tmap!, SerialScheduler
 using DocStringExtensions
 
 include("Defaults.jl")  # Include first to allow for docstring interpolation with Defaults values
diff --git a/src/algorithms/contractions/absorb_weight.jl b/src/algorithms/contractions/absorb_weight.jl
index 6362b1aca..8f8ce451e 100644
--- a/src/algorithms/contractions/absorb_weight.jl
+++ b/src/algorithms/contractions/absorb_weight.jl
@@ -48,41 +48,52 @@ function absorb_weight(
         t::Union{PEPSTensor, PEPOTensor}, weights::SUWeight,
         row::Int, col::Int, virt_axes::NTuple{N, Int}; inv::Bool = false
     ) where {N}
-    return _absorb_weight(t, weights, row, col, virt_axes; inv)
+    wts = map(ax -> weight_to_absorb(weights, row, col, ax; inv), virt_axes)
+    T = promote_type(scalartype(t), map(scalartype, wts)...)
+    return _absorb_weights!(similar(t, T), t, wts, virt_axes)
 end
 
 """
-Absorb weights on the virtual directions in `virt_axes`, or remove them when `inv = true`.
+    absorb_weight!(t, weights::SUWeight, rowcol::CartesianIndex{2}, virt_axes; inv = false)
 
-- `virt_axes` selects virtual directions using `1 = NORTH`, `2 = EAST`, etc., regardless of the current leg order.
-- `input_axes` specifies how the legs of the input `t` are already permuted with respect to the standard order.
-  For example, `input_axes = (5, 1, 2, 4, 3)` for a PEPS tensor means the legs are currently ordered as west, physical, north, south, east.
-- `output_perm` specifies how the returned tensor should permute the legs with respect to the standard order.
-  For example, `output_perm = ((5, 1, 2, 4), (3,))` puts west, physical, north, south in the codomain and east in the domain.
-- There are always four virtual legs on the square lattice, so the physical leg count is `numind(t) - 4`.
+In-place version of [`absorb_weight`](@ref).
 """
-function _absorb_weight(
-        t::AbstractTensorMap, weights::SUWeight, row::Int, col::Int,
-        virt_axes::NTuple{N, Int}; inv::Bool = false,
-        input_axes::Tuple = ntuple(identity, numind(t)),
-        output_perm::Tuple{Tuple, Tuple} = (
-            ntuple(identity, numind(t) - 4), ntuple(i -> numind(t) - 4 + i, 4),
-        )
+function absorb_weight!(
+        t::Union{PEPSTensor, PEPOTensor}, weights::SUWeight,
+        rowcol::CartesianIndex{2}, virt_axes::NTuple{N, Int}; inv::Bool = false
+    ) where {N}
+    wts = map(ax -> weight_to_absorb(weights, rowcol[1], rowcol[2], ax; inv), virt_axes)
+    return _absorb_weights!(t, t, wts, virt_axes)
+end
+
+# Write `src` into `dst` with `weights[i]` absorbed into virtual axis `virtual_axes[i]`,
+# where `domain(weights[i]) == domain(src)[virtual_axes[i]]`; `dst === src` is allowed.
+# A diagonal weight only rescales the basis vectors of its leg, by a factor that depends on
+# the sector of the leg and the index within it. This commutes with any permutation, so it is
+# applied directly on each fusion-tree block of `src` in its own leg order; it is equivalent
+# to `permute(src, ...) * wt` followed by the inverse permutation.
+# Since the domain of `wt` is the space of the leg itself (dual legs get `transpose(wt)` from
+# `weight_to_absorb`), `wt` is composed onto the leg without bending it, so even for dual legs
+# no twist is involved.
+function _absorb_weights!(
+        dst::AbstractTensorMap, src::AbstractTensorMap,
+        weights::NTuple{N, DiagonalTensorMap}, virtual_axes::NTuple{N, Int}
     ) where {N}
-    Np = numind(t) - 4
-    vax = first(virt_axes)
-    axes, biperm = _permute_to_last(input_axes, vax + Np)
-    wt = weight_to_absorb(weights, row, col, vax; inv)
-    # Allocate the first result to preserve the input and allow scalar promotion.
-    t2 = permute(t, biperm) * wt
-    for vax in Base.tail(virt_axes)
-        axes, biperm = _permute_to_last(axes, vax + Np)
-        wt = weight_to_absorb(weights, row, col, vax; inv)
-        # Diagonal composition is twist-free and only mutates the intermediate.
-        t2 = rmul!(permute(t2, biperm), wt)
+    Np, M = numout(src), numind(src)
+    space(dst) == space(src) ||
+        throw(SpaceMismatch(lazy"destination $(space(dst)) does not match source $(space(src))"))
+    for (wt, ax) in zip(weights, virtual_axes)
+        domain(wt)[1] == domain(src)[ax] ||
+            throw(SpaceMismatch(lazy"weight $(space(wt)) does not match virtual axis $ax of $(space(src))"))
+    end
+    for (f₁, f₂) in fusiontrees(src)
+        λs = map(weights, virtual_axes) do wt, ax
+            λ = StridedView(block(wt, f₂.uncoupled[ax]).diag)
+            return sreshape(λ, ntuple(i -> i == Np + ax ? length(λ) : 1, M))
+        end
+        @. dst[f₁, f₂] = *(src[f₁, f₂], λs...)
     end
-    perm_back = invperm(axes)
-    return permute(t2, map(p -> map(i -> perm_back[i], p), output_perm))
+    return dst
 end
 
 """
diff --git a/src/algorithms/time_evolution/apply_gate.jl b/src/algorithms/time_evolution/apply_gate.jl
index 9b6f811aa..5ed1c4cf8 100644
--- a/src/algorithms/time_evolution/apply_gate.jl
+++ b/src/algorithms/time_evolution/apply_gate.jl
@@ -20,33 +20,112 @@ function _apply_sitegate(
     return a′
 end
 
+# --- 2-site gates on nearest neighbor bonds ---------------------------------------------
+#
+# The bond between `s₁` and `s₂` is updated in the native PEPS/PEPO leg order, using the
+# bond direction to locate the legs. Each site tensor is split as `T = Q * R`, where the
+# reduced tensor `R` holds the physical leg acted on by the gate and the bond leg. The gate is
+# applied to `R₁ -- R₂` and split again by a truncated SVD, after which `Q` and the new `R`
+# are contracted back into the native leg order.
+#
+# Temporaries allocated with `allocator` never outlive the function that allocates them:
+# each function below frees its own temporaries and resets `allocator` before returning.
+
+"""
+    _orient_bond(state, sites, gate) -> (s₁, k₁), (s₂, k₂), gate
+
+Order the two `sites` of a nearest neighbor bond and the virtual axes `k₁`, `k₂` of their
+bond legs, such that the bond leg of `s₁` is non-dual. Putting `s₁` on the codomain side
+of the SVD then reproduces the original arrow on the bond, so nothing has to be flipped.
+"""
+function _orient_bond(state::InfiniteState, sites::Vector{CartesianIndex{2}}, gate::NNGate)
+    s₁, s₂ = sites
+    k₁ = first(_nn_bonddir(s₁, s₂))
+    k₂ = mod1(k₁ + 2, 4)
+    T₁ = state[s₁]
+    # `space` of a domain leg returns its dual
+    isdual(space(T₁, numout(T₁) + k₁)) && return (s₁, k₁), (s₂, k₂), gate
+    return (s₂, k₂), (s₁, k₁), permute(gate, ((2, 1), (4, 3)))
+end
+
 """
-$(SIGNATURES)
+    _bond_qr(T, k, gate_ax, [allocator]) -> Q, R, p
 
-Apply 2-site `gate` on the reduced bond tensors `a`, `b`
+Split `T = Q * R`, where the reduced tensor `R : x ← p ⊗ v` holds the physical axis `gate_ax`
+and the virtual axis `k` (the bond) of `T`, while `Q` holds all other axes. `p` is the
+permutation that brings the product back to the leg order of `T`. For example, for a
+PEPSTensor `T[p; n e s w]` with the bond `v = e` on its east axis (`k = 2`),
 ```
-    -1← a --- 3 --- b ← -4          -2         -3
-        ↓           ↓               ↓           ↓
-        1           2               |----gate---|
-        ↓           ↓       or      ↓           ↓
-        |----gate---|               1           2
-        ↓           ↓               ↓           ↓
-        -2         -3           -1← a --- 3 --- b ← -4
+        n                       n
+        |                       |
+    w - T - v       =       w - Q ← x ← R - v
+        | ↘                     |        ↘
+        s   p                   s         p
 ```
+For a PEPOTensor, the physical axis that is not `gate_ax` stays in `Q`.
+The permuted copy of `T` that is factorized is a temporary allocated with `allocator`.
 """
-function _apply_gate(a::MPSTensor, b::MPSTensor, gate::NNGate, trunc::TruncationStrategy)
-    V = space(b, 1)
-    need_flip = isdual(V)
-    if isdual(space(a, 2))
-        @tensor a2b2[-1 -2; -3 -4] := gate[1 2; -2 -3] * a[-1 1 3] * b[3 2 -4]
+function _bond_qr(
+        T::Union{PEPSTensor, PEPOTensor}, k::Int, gate_ax::Int, allocator = DefaultAllocator()
+    )
+    Np, M = numout(T), numind(T)
+    r = (gate_ax, Np + k)
+    q = TupleTools.deleteat(ntuple(identity, M), r)
+    cp = allocator_checkpoint!(allocator)
+    T′ = tensoralloc_add(scalartype(T), T, (q, r), false, Val(true), allocator)
+    tensoradd!(T′, T, (q, r), false, One(), Zero(), DefaultBackend(), allocator)
+    Q, R = left_orth!(T′; positive = true)
+    tensorfree!(T′, allocator)
+    allocator_reset!(allocator, cp)
+    return Q, R, _inv_mpo_perm((q..., r...), Val(Np))
+end
+
+"""
+    _undo_bond_qr(Q, X, pX, p, [allocator])
+
+Inverse of [`_bond_qr`](@ref): contract `Q` with `X`, whose legs selected by `pX` are
+ordered as `x ← p ⊗ v`, into the leg order specified by `p`.
+"""
+function _undo_bond_qr(
+        Q::AbstractTensorMap, X::AbstractTensorMap, pX, p, allocator = DefaultAllocator()
+    )
+    N = numind(Q)
+    pQ = (ntuple(identity, N - 1), (N,))
+    cp = allocator_checkpoint!(allocator)
+    T = tensorcontract(Q, pQ, false, X, pX, false, p, One(), DefaultBackend(), allocator)
+    allocator_reset!(allocator, cp)
+    return T
+end
+
+"""
+    _apply_bond_gate(R₁, R₂, gate, trunc, [allocator]) -> U, S, Vᴴ, ϵ
+
+Apply the 2-site `gate` to the reduced tensors `R₁ : x₁ ← p₁ ⊗ v` and `R₂ : x₂ ← p₂ ⊗ v'`
+```
+    x₁ ← R₁ --- v --- R₂ → x₂           p₁            p₂
+         ↘            ↘                 ↓             ↓
+          p₁           p₂               |----gate-----|
+          ↓            ↓        or      ↓             ↓
+          |----gate----|                p₁            p₂
+          ↓            ↓                ↘             ↘
+          p₁           p₂          x₁ ← R₁ --- v --- R₂ → x₂
+```
+where the second form is used when the physical legs are dual (the second physical axis of
+a PEPO), and split the result as `U S Vᴴ` with a truncated SVD. The square root of `S` is
+absorbed into `U : x₁ ⊗ p₁ ← v′` and `Vᴴ : v′ ← x₂ ⊗ p₂`.
+"""
+function _apply_bond_gate(
+        R₁::AbstractTensorMap, R₂::AbstractTensorMap, gate::NNGate,
+        trunc::TruncationStrategy, allocator = DefaultAllocator()
+    )
+    θ = if isdual(space(R₁, 2))
+        @tensor allocator = allocator θ[-1 -2; -3 -4] :=
+            R₁[-1; 1 3] * R₂[-3; 2 3] * gate[1 2; -2 -4]
     else
-        @tensor a2b2[-1 -2; -3 -4] := gate[-2 -3; 1 2] * a[-1 1 3] * b[3 2 -4]
-    end
-    a, s, b, ϵ = svd_trunc!(a2b2; trunc)
-    a, b = absorb_s(a, s, b)
-    if need_flip
-        a, s, b = flip(a, numind(a)), _fliptwist_s!(s), flip(b, 1)
+        @tensor allocator = allocator θ[-1 -2; -3 -4] :=
+            R₁[-1; 1 3] * R₂[-3; 2 3] * gate[-2 -4; 1 2]
     end
-    b = permute(b, ((1, 2), (3,)))
-    return a, s, b, ϵ
+    U, S, Vᴴ, ϵ = svd_trunc!(θ; trunc)
+    absorb_s!(U, S, Vᴴ)
+    return U, S, Vᴴ, ϵ
 end
diff --git a/src/algorithms/time_evolution/get_cluster.jl b/src/algorithms/time_evolution/get_cluster.jl
index c9a97c629..da94c9449 100644
--- a/src/algorithms/time_evolution/get_cluster.jl
+++ b/src/algorithms/time_evolution/get_cluster.jl
@@ -72,22 +72,6 @@ function _nn_bonddir(site1::CartesianIndex{2}, site2::CartesianIndex{2})
     end
 end
 
-"""
-Apply the tensor rotation that maps a bond from orientation `dir_from` to orientation `dir_to`, where the directions are one of `NORTH`, `EAST`, `SOUTH`, `WEST`.
-"""
-function _bond_rotation(x, dir_from::Int, dir_to::Int)
-    delta = mod(dir_to - dir_from, 4)
-    return if delta == 0
-        x
-    elseif delta == 1
-        rotr90(x)
-    elseif delta == 2
-        rot180(x)
-    else # delta == 3
-        rotl90(x)
-    end
-end
-
 """
 Obtain the cluster `Ms` along the (open) path `sites` in `state`.
 
@@ -136,27 +120,3 @@ function _get_cluster(
     Ms = map(site -> state[site], sites)
     return Ms, open_vaxs, perms
 end
-
-"""
-Absorb weights into the open virtual legs of
-the cluster `Ms` of PEPS/PEPO tensors.
-"""
-function _absorb_weight!(
-        Ms::Vector{T}, sites::Vector{CartesianIndex{2}},
-        open_vaxs::Vector{<:Tuple}, env::SUWeight; inv::Bool = false
-    ) where {T}
-    for (idx, (M, site, vaxs)) in enumerate(zip(Ms, sites, open_vaxs))
-        Ms[idx] = absorb_weight(M, env, site, vaxs; inv)
-    end
-    return Ms
-end
-
-"""
-Given a vector `Ms` of `AbstractTensorMap`s,
-apply permutations `perms` to each tensor.
-"""
-function _permute_cluster(Ms::Vector{<:AbstractTensorMap}, perms::Vector{<:Tuple})
-    return map(Ms, perms) do M, perm
-        return permute(M, perm)
-    end
-end
diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index fb64c0f47..38a3614ea 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -76,43 +76,44 @@ end
 """
 Simple update optimized for nearest neighbor gates
 utilizing reduced bond tensors with the physical leg.
+
+The tensors are updated in their native leg order, without rotating them.
+Temporary tensors are allocated with `allocator`.
 """
 function _su_iter!(
         state::InfiniteState, gate::NNGate, env::SUWeight,
-        sites::Vector{CartesianIndex{2}}, alg::SimpleUpdate
+        sites::Vector{CartesianIndex{2}}, alg::SimpleUpdate;
+        allocator = DefaultAllocator()
     )
-    @assert length(sites) == 2
-    trunc = only(_get_cluster_trunc(alg.trunc, sites))
-    Ms, open_vaxs, = _get_cluster(state, sites)
-    _absorb_weight!(Ms, sites, open_vaxs, env)
-    # rotate
-    (dir, r, c) = _nn_bonddir(sites...)
-    A, B = _bond_rotation.(Ms, dir, EAST)
+    (s₁, k₁), (s₂, k₂), gate = _orient_bond(state, sites, gate)
+    trunc = only(_get_cluster_trunc(alg.trunc, [s₁, s₂]))
+    open₁ = TupleTools.deleteat((1, 2, 3, 4), k₁)
+    open₂ = TupleTools.deleteat((1, 2, 3, 4), k₂)
+    # absorb environment weights
+    T₁ = absorb_weight(state[s₁], env, s₁, open₁)
+    T₂ = absorb_weight(state[s₂], env, s₂, open₂)
     # apply gate
     ϵ = 0.0
-    local s
-    for gate_ax in 1:2
-        a, X = bond_tensor_first(A; gate_ax, positive = true)
-        b, Y = bond_tensor_last(B; gate_ax, positive = true)
-        a, s, b, ϵ′ = _apply_gate(a, b, gate, trunc)
+    local S
+    for gate_ax in 1:numout(T₁)
+        Q₁, R₁, p₁ = _bond_qr(T₁, k₁, gate_ax, allocator)
+        Q₂, R₂, p₂ = _bond_qr(T₂, k₂, gate_ax, allocator)
+        U, S, Vᴴ, ϵ′ = _apply_bond_gate(R₁, R₂, gate, trunc, allocator)
         ϵ = max(ϵ, ϵ′)
-        A = undo_bond_tensor_first(a, X; gate_ax)
-        B = undo_bond_tensor_last(b, Y; gate_ax)
+        T₁ = _undo_bond_qr(Q₁, U, ((1,), (2, 3)), p₁, allocator)
+        T₂ = _undo_bond_qr(Q₂, Vᴴ, ((2,), (3, 1)), p₂, allocator)
         alg.purified && break # only apply gate to 1st physical leg
     end
-    # rotate back
-    A = _bond_rotation(A, EAST, dir)
-    B = _bond_rotation(B, EAST, dir)
-    (dir in (WEST, SOUTH)) && (s = transpose(s))
     # remove environment weights
-    siteA, siteB = sites
-    A = absorb_weight(A, env, siteA[1], siteA[2], open_vaxs[1]; inv = true)
-    B = absorb_weight(B, env, siteB[1], siteB[2], open_vaxs[2]; inv = true)
-    # update tensor dict and weight on current bond
-    state[siteA] = normalize!(A, Inf)
-    state[siteB] = normalize!(B, Inf)
+    absorb_weight!(T₁, env, s₁, open₁; inv = true)
+    absorb_weight!(T₂, env, s₂, open₂; inv = true)
+    # update tensor dict and weight on current bond,
+    # which lives on the EAST (NORTH) leg of the west (south) site
+    state[s₁] = normalize!(T₁, Inf)
+    state[s₂] = normalize!(T₂, Inf)
+    dir, r, c = _nn_bonddir(s₁, s₂)
     d = dir in (EAST, WEST) ? 1 : 2
-    env[d, r, c] = normalize!(s, Inf)
+    env[d, r, c] = normalize!(dir in (EAST, NORTH) ? S : transpose(S), Inf)
     return ϵ
 end
 
@@ -124,6 +125,8 @@ function su_iter(
         alg::SimpleUpdate, env::SUWeight
     )
     state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0
+    # the gates are applied one after another, so a single allocator serves all of them
+    allocator = default_allocator(state, SerialScheduler())
     for (sites, gate) in circuit.gates
         if length(sites) == 1
             # 1-site gate
@@ -134,7 +137,7 @@ function su_iter(
             (dir, r, c) = _nn_bonddir(sites...)
             d = dir in (EAST, WEST) ? 1 : 2
             alg.bipartite && iseven(r) && continue
-            ϵ′ = _su_iter!(state2, gate, env2, sites, alg)
+            ϵ′ = _su_iter!(state2, gate, env2, sites, alg; allocator)
             ϵ = max(ϵ, ϵ′)
             (!alg.bipartite) && continue
             if d == 1
diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl
index a0f55c693..1614f6e79 100644
--- a/src/algorithms/time_evolution/simpleupdate3site.jl
+++ b/src/algorithms/time_evolution/simpleupdate3site.jl
@@ -14,7 +14,8 @@ function _unfuse_physicalspace(
 end
 
 """
-Absorb open-leg weights directly into the cluster's MPS layout.
+Obtain the cluster along `sites` with weights absorbed into its open virtual legs,
+and permuted to MPS axis order.
 """
 function _get_cluster_with_weights(
         state::InfiniteState, sites::Vector{CartesianIndex{2}}, env::SUWeight
@@ -22,7 +23,7 @@ function _get_cluster_with_weights(
     Ms, open_vaxs, perms = _get_cluster(state, sites)
     Np = (state isa InfinitePEPS) ? Val(1) : Val(2)
     Ms_mps = map(Ms, sites, open_vaxs, perms) do M, site, vaxs, perm
-        return _absorb_weight(M, env, site[1], site[2], vaxs; output_perm = perm)
+        return permute(absorb_weight(M, env, site, vaxs), perm)
     end
     invperms = map(p -> _inv_mpo_perm(p, Np), perms)
     return Ms_mps, open_vaxs, invperms
@@ -30,10 +31,14 @@ end
 
 """
 Simple update with an N-site MPO `gate` (N ≥ 2).
+
+The `allocator` is accepted for compatibility with the nearest neighbor update,
+but is not used yet.
 """
 function _su_iter!(
         state::InfiniteState, gate::Vector{T}, env::SUWeight,
-        sites::Vector{CartesianIndex{2}}, alg::SimpleUpdate
+        sites::Vector{CartesianIndex{2}}, alg::SimpleUpdate;
+        allocator = DefaultAllocator()
     ) where {T <: AbstractTensorMap}
     truncs = _get_cluster_trunc(alg.trunc, sites)
     Ms, open_vaxs, invperms = _get_cluster_with_weights(state, sites, env)
@@ -62,10 +67,9 @@ function _su_iter!(
         env[d, r, c] = normalize!(wt_new, Inf)
     end
     # update state tensors
-    for (M, s, perm_back, vaxs) in zip(Ms, sites, invperms, open_vaxs)
-        # Remove weights directly from MPS order, then restore the site's layout.
-        input_axes = invperm((perm_back[1]..., perm_back[2]...))
-        M′ = _absorb_weight(M, env, s[1], s[2], vaxs; inv = true, input_axes)
+    for (M, s, invperm, vaxs) in zip(Ms, sites, invperms, open_vaxs)
+        # restore original axes order and remove weights on open axes of the cluster
+        M′ = absorb_weight!(permute(M, invperm), env, s, vaxs; inv = true)
         state[s] = normalize!(M′, Inf)
     end
     return ϵ
diff --git a/src/utility/util.jl b/src/utility/util.jl
index 9893262d3..3f8960e7a 100644
--- a/src/utility/util.jl
+++ b/src/utility/util.jl
@@ -66,17 +66,24 @@ end
 
 """
     absorb_s(U::AbstractTensorMap, S::DiagonalTensorMap, V::AbstractTensorMap)
+    absorb_s!(U::AbstractTensorMap, S::DiagonalTensorMap, V::AbstractTensorMap)
 
 Given SVD result `U`, `S` and `V`, absorb singular values `S` into `U` and `V` by:
 ```
     U -> U * sqrt(S), V -> sqrt(S) * V
 ```
+The in-place version `absorb_s!` overwrites `U` and `V`.
 """
 function absorb_s(U::AbstractTensorMap, S::DiagonalTensorMap, V::AbstractTensorMap)
     @assert !isdual(space(S, 1))
     sqrt_S = sdiag_pow(S, 0.5)
     return U * sqrt_S, sqrt_S * V
 end
+function absorb_s!(U::AbstractTensorMap, S::DiagonalTensorMap, V::AbstractTensorMap)
+    @assert !isdual(space(S, 1))
+    sqrt_S = sdiag_pow(S, 0.5)
+    return rmul!(U, sqrt_S), lmul!(sqrt_S, V)
+end
 
 # avoids a conversion to a dense TensorMap by flip
 function _fliptwist_s!(s::DiagonalTensorMap)
@@ -168,14 +175,3 @@ function random_dual!(Vs::AbstractMatrix{E}; p = 0.7) where {E <: ElementarySpac
     return Vs
 end
 
-"""
-Returns `(1, 2, ..., N)` but with `ax` moved to the end,
-and the corresponding permutation for `axes` (with `ax` as the only domain index).
-"""
-function _permute_to_last(axes::NTuple{N, Int}, ax::Int) where {N}
-    codomain_axes = TupleTools.deleteat(ntuple(identity, N), ax)
-    q = invperm(axes)
-    biperm = (map(i -> q[i], codomain_axes), (q[ax],))
-    new_axes = (codomain_axes..., ax)
-    return new_axes, biperm
-end

From 180d0740ed160f25ed0bfeff88e7089940198309 Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Wed, 28 Oct 2026 18:12:48 -0400
Subject: [PATCH 5/7] Add simple update benchmarks

Add a `timeevol` benchmark module that times one `su_iter` sweep over a
2 x 2 unit cell for trivial and U(1) Heisenberg PEPS, fermionic Hubbard
PEPS and a Heisenberg PEPO, at two bond dimensions each.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../PEPSKitBenchmarks/PEPSKitBenchmarks.jl    |  2 +
 .../timeevol/TimeEvolBenchmarks.jl            | 76 +++++++++++++++++++
 benchmark/README.md                           | 34 +++++++--
 3 files changed, 104 insertions(+), 8 deletions(-)
 create mode 100644 benchmark/PEPSKitBenchmarks/timeevol/TimeEvolBenchmarks.jl

diff --git a/benchmark/PEPSKitBenchmarks/PEPSKitBenchmarks.jl b/benchmark/PEPSKitBenchmarks/PEPSKitBenchmarks.jl
index 2d41d668e..a33ae4f92 100644
--- a/benchmark/PEPSKitBenchmarks/PEPSKitBenchmarks.jl
+++ b/benchmark/PEPSKitBenchmarks/PEPSKitBenchmarks.jl
@@ -14,9 +14,11 @@ const PARAMS_PATH = joinpath(@__DIR__, "etc", "params.json")
 const SUITE = BenchmarkGroup()
 const MODULES = Dict{String, Symbol}(
     "ctmrg" => :CTMRGBenchmarks,
+    "timeevol" => :TimeEvolBenchmarks,
 )
 
 include("ctmrg/CTMRGBenchmarks.jl")
+include("timeevol/TimeEvolBenchmarks.jl")
 
 load!(id::AbstractString; kwargs...) = load!(SUITE, id; kwargs...)
 
diff --git a/benchmark/PEPSKitBenchmarks/timeevol/TimeEvolBenchmarks.jl b/benchmark/PEPSKitBenchmarks/timeevol/TimeEvolBenchmarks.jl
new file mode 100644
index 000000000..4b2cefd5b
--- /dev/null
+++ b/benchmark/PEPSKitBenchmarks/timeevol/TimeEvolBenchmarks.jl
@@ -0,0 +1,76 @@
+module TimeEvolBenchmarks
+
+using BenchmarkTools
+using TensorKit
+using PEPSKit
+using PEPSKit.MPSKitModels: hubbard_space
+
+const SUITE = BenchmarkGroup()
+
+"""
+A simple update scenario on a 2 x 2 unit cell: the Hamiltonian `H` and the initial state `psi`,
+evolved with bond dimension `D` (the truncation keeps the initial virtual spaces).
+"""
+struct SUSpec{H, S}
+    H::H
+    psi::S
+    D::Int
+    purified::Bool
+end
+
+function heisenberg_peps(D::Int; T = Float64)
+    H = heisenberg_XYZ(InfiniteSquare(2, 2))
+    psi = InfinitePEPS(randn, T, ℂ^2, ℂ^D; unitcell = (2, 2))
+    return SUSpec(H, psi, D, true)
+end
+
+function heisenberg_u1_peps(n::Int; T = Float64)
+    H = heisenberg_XXZ(T, U1Irrep, InfiniteSquare(2, 2))
+    V = U1Space(0 => 2n, 1 // 2 => n, -1 // 2 => n)
+    psi = InfinitePEPS(randn, T, physicalspace(H)[1, 1], V; unitcell = (2, 2))
+    return SUSpec(H, psi, dim(V), true)
+end
+
+function hubbard_peps(n::Int; T = Float64)
+    H = hubbard_model(T, Trivial, U1Irrep, InfiniteSquare(2, 2); t = 1.0, U = 6.0, mu = 3.0)
+    V = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 2n, (1, 1 // 2) => n, (1, -1 // 2) => n)
+    psi = InfinitePEPS(randn, T, hubbard_space(Trivial, U1Irrep), V; unitcell = (2, 2))
+    return SUSpec(H, psi, dim(V), true)
+end
+
+function heisenberg_pepo(D::Int; T = Float64)
+    H = heisenberg_XYZ(InfiniteSquare(2, 2))
+    psi = InfinitePEPO(randn, T, ℂ^2, ℂ^D; unitcell = (2, 2, 1))
+    return SUSpec(H, psi, D, false)
+end
+
+"""
+Benchmark one simple update iteration, i.e. the application of all Trotter gates of `spec.H`,
+starting from weights obtained from a few warm-up iterations.
+"""
+function su_iter_benchmark(spec::SUSpec; dt = 1.0e-2, nwarmup = 3)
+    alg = SimpleUpdate(; trunc = truncrank(spec.D), spec.purified)
+    circuit = PEPSKit.trotterize(spec.H, dt)
+    psi, env = spec.psi, SUWeight(spec.psi)
+    for _ in 1:nwarmup
+        psi, env, = PEPSKit.su_iter(psi, circuit, alg, env)
+    end
+    return @benchmarkable PEPSKit.su_iter($psi, $circuit, $alg, $env)
+end
+
+const SCENARIOS = [
+    ("Heisenberg PEPS", "Trivial", "D8") => () -> heisenberg_peps(8),
+    ("Heisenberg PEPS", "Trivial", "D16") => () -> heisenberg_peps(16),
+    ("Heisenberg PEPS", "U1", "D8") => () -> heisenberg_u1_peps(2),
+    ("Heisenberg PEPS", "U1", "D16") => () -> heisenberg_u1_peps(4),
+    ("Hubbard PEPS", "fZ2xU1", "D8") => () -> hubbard_peps(2),
+    ("Hubbard PEPS", "fZ2xU1", "D16") => () -> hubbard_peps(4),
+    ("Heisenberg PEPO", "Trivial", "D4") => () -> heisenberg_pepo(4),
+    ("Heisenberg PEPO", "Trivial", "D6") => () -> heisenberg_pepo(6),
+]
+
+for (groupkeys, spec) in SCENARIOS
+    SUITE[("su_iter", groupkeys...)] = su_iter_benchmark(spec())
+end
+
+end # module
diff --git a/benchmark/README.md b/benchmark/README.md
index 9a9290c4a..f530291d9 100644
--- a/benchmark/README.md
+++ b/benchmark/README.md
@@ -1,9 +1,13 @@
 # PEPSKit benchmarks
 
 A [BenchmarkTools.jl](https://github.com/JuliaCI/BenchmarkTools.jl)-based suite that exercises
-performance-critical paths in PEPSKit. The suite is organised into modules; the only module
-currently registered is `ctmrg`, which times a single `ctmrg_iteration` across several CTMRG
-algorithms, problem scenarios, and symmetry sectors.
+performance-critical paths in PEPSKit. The suite is organised into modules:
+
+- `ctmrg` times a single `ctmrg_iteration` across several CTMRG algorithms, problem scenarios,
+  and symmetry sectors.
+- `timeevol` times a single simple update iteration (`su_iter`, i.e. one sweep of all Trotter
+  gates over a 2 x 2 unit cell) for PEPS and PEPO time evolution, across models, symmetry sectors
+  and bond dimensions.
 
 ## Layout
 
@@ -16,11 +20,13 @@ bench/benchmark/
 └── PEPSKitBenchmarks/
     ├── PEPSKitBenchmarks.jl      # module registry (MODULES), load!, loadall!
     ├── utils/BenchUtils.jl       # TOML <-> VectorSpace conversion helpers
-    └── ctmrg/
-        ├── CTMRGBenchmarks.jl    # populates SUITE["ctmrg"]
-        ├── ctmrg_iteration_benchmarks.jl
-        ├── default.toml          # Trivial + NonUniform scenarios
-        └── su3_hubbard.toml      # SU(3) fermionic scenarios
+    ├── ctmrg/
+    │   ├── CTMRGBenchmarks.jl    # populates SUITE["ctmrg"]
+    │   ├── ctmrg_iteration_benchmarks.jl
+    │   ├── default.toml          # Trivial + NonUniform scenarios
+    │   └── su3_hubbard.toml      # SU(3) fermionic scenarios
+    └── timeevol/
+        └── TimeEvolBenchmarks.jl # populates SUITE["timeevol"]
 ```
 
 New modules are registered in the `MODULES` dict in `PEPSKitBenchmarks.jl`.
@@ -54,6 +60,18 @@ julia --project=bench/benchmark bench/benchmark/benchmarks.jl --modules=ctmrg
 The script only populates `SUITE`; call `run(SUITE)` afterwards (e.g. from an interactive
 session) to execute. Available IDs are the keys of `MODULES` in `PEPSKitBenchmarks.jl`.
 
+### Simple update
+
+`SUITE["timeevol"]` is indexed as `"su_iter" → model → symmetry → "D<D>"`, e.g.
+
+```julia
+include("bench/benchmark/benchmarks.jl")
+run(SUITE["timeevol"]["su_iter"]["Hubbard PEPS"]["fZ2xU1"]["D16"])
+```
+
+The scenarios are defined in the `SCENARIOS` list in `timeevol/TimeEvolBenchmarks.jl`; each
+entry maps its group keys to a function that builds the Hamiltonian and initial state.
+
 ### By group key
 
 `SUITE` is a nested `BenchmarkGroup` indexed as:

From 0e1c18e1617e580d0fee4b69016bcb38743c3e28 Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Wed, 28 Oct 2026 18:24:00 -0400
Subject: [PATCH 6/7] Hoist structure lookups out of the weight absorption loop
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Per fusion-tree block, `block(wt, c)` and `t[f₁, f₂]` each redo a lookup in
TensorKit's global structure cache, which dominated the cost of absorbing
weights into tensors with many small blocks. Look up the weight blocks once
per call and iterate with `subblocks` instead. Also make the tensor rank
static in the `ntuple` that reshapes the weights, which was otherwise not
inferred and caused dynamic dispatch on every broadcast.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/algorithms/contractions/absorb_weight.jl | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/algorithms/contractions/absorb_weight.jl b/src/algorithms/contractions/absorb_weight.jl
index 8f8ce451e..79b27c7a0 100644
--- a/src/algorithms/contractions/absorb_weight.jl
+++ b/src/algorithms/contractions/absorb_weight.jl
@@ -79,19 +79,25 @@ function _absorb_weights!(
         dst::AbstractTensorMap, src::AbstractTensorMap,
         weights::NTuple{N, DiagonalTensorMap}, virtual_axes::NTuple{N, Int}
     ) where {N}
-    Np, M = numout(src), numind(src)
+    Np, M = numout(src), Val(numind(src)) # static rank, so the reshaped weights are inferred
     space(dst) == space(src) ||
         throw(SpaceMismatch(lazy"destination $(space(dst)) does not match source $(space(src))"))
     for (wt, ax) in zip(weights, virtual_axes)
         domain(wt)[1] == domain(src)[ax] ||
             throw(SpaceMismatch(lazy"weight $(space(wt)) does not match virtual axis $ax of $(space(src))"))
     end
-    for (f₁, f₂) in fusiontrees(src)
-        λs = map(weights, virtual_axes) do wt, ax
-            λ = StridedView(block(wt, f₂.uncoupled[ax]).diag)
-            return sreshape(λ, ntuple(i -> i == Np + ax ? length(λ) : 1, M))
-        end
-        @. dst[f₁, f₂] = *(src[f₁, f₂], λs...)
+    # Look up every block of the weights only once, reshaped to broadcast along their leg, and
+    # iterate over the subblocks of `dst` and `src` (which have the same order) with `subblocks`,
+    # as the per-block `block(wt, c)` and `t[f₁, f₂]` would both redo the structure lookups.
+    λdicts = map(weights, virtual_axes) do wt, ax
+        return Dict(
+            c => sreshape(StridedView(b.diag), ntuple(i -> i == Np + ax ? length(b.diag) : 1, M))
+                for (c, b) in blocks(wt)
+        )
+    end
+    for (((f₁, f₂), d), (_, s)) in zip(subblocks(dst), subblocks(src))
+        λs = map((λdict, ax) -> λdict[f₂.uncoupled[ax]], λdicts, virtual_axes)
+        @. d = *(s, λs...)
     end
     return dst
 end

From d13a8d223c149db58cd73837c2b5b6c2ff4e4457 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Thu, 29 Oct 2026 09:26:26 +0800
Subject: [PATCH 7/7] Formatting

---
 src/utility/util.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/utility/util.jl b/src/utility/util.jl
index 3f8960e7a..df25d2ae5 100644
--- a/src/utility/util.jl
+++ b/src/utility/util.jl
@@ -174,4 +174,3 @@ function random_dual!(Vs::AbstractMatrix{E}; p = 0.7) where {E <: ElementarySpac
     end
     return Vs
 end
-