Add bipermutedimsopadd! overloads

Project.toml

@@ -1,6 +1,6 @@
 name = "GradedArrays"
 uuid = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
-version = "0.8.3"
+version = "0.8.4"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -48,7 +48,7 @@ Random = "1.10"
 SUNRepresentations = "0.3"
 SparseArraysBase = "0.9"
 SplitApplyCombine = "1.2.3"
-TensorAlgebra = "0.8"
+TensorAlgebra = "0.8, 0.9"
 TensorKitSectors = "0.3"
 TypeParameterAccessors = "0.4"
 julia = "1.10"

src/GradedArrays.jl

@@ -30,9 +30,9 @@ using BlockSparseArrays: BlockSparseArrays, blockdiagindices, blockstoredlength,
 using KroneckerArrays: KroneckerArrays, kroneckerfactors, ×, ⊗
 using LinearAlgebra: LinearAlgebra, Adjoint, mul!
 using SparseArraysBase: SparseArraysBase
-using TensorAlgebra: TensorAlgebra, BlockedTuple, FusionStyle, matricize, matricize_axes,
-    permutedimsadd!, permutedimsopadd!, tensor_product_axis, trivial_axis, trivialbiperm,
-    tryflattenlinear, unmatricize
+using TensorAlgebra: TensorAlgebra, BlockedTuple, FusionStyle, bipermutedimsopadd!,
+    check_input, matricize, matricize_axes, permutedimsadd!, tensor_product_axis,
+    trivial_axis, trivialbiperm, tryflattenlinear, unmatricize
 using TensorKitSectors: TensorKitSectors as TKS
 using TypeParameterAccessors: type_parameters, unspecify_type_parameters
 

src/fusedgradedmatrix.jl

@@ -143,13 +143,13 @@ end
 
 # ========================  mul!  ========================
 
-function check_input(::typeof(*), A::FusedGradedMatrix, B::FusedGradedMatrix)
+function TensorAlgebra.check_input(::typeof(*), A::FusedGradedMatrix, B::FusedGradedMatrix)
     axes(A, 2) == dual(axes(B, 1)) ||
         throw(DimensionMismatch("sector mismatch in contracted dimension"))
     return nothing
 end
 
-function check_input(
+function TensorAlgebra.check_input(
         ::typeof(mul!),
         C::FusedGradedMatrix, A::FusedGradedMatrix, B::FusedGradedMatrix
     )

src/fusion.jl

@@ -158,7 +158,7 @@ end
 # (via `sort!(vcat(...))`) would reorder blocks unexpectedly otherwise.
 function insert_missing_sectors end
 
-function check_input(::typeof(insert_missing_sectors), a::AbelianGradedMatrix)
+function TensorAlgebra.check_input(::typeof(insert_missing_sectors), a::AbelianGradedMatrix)
     cod_sects = sectors(axes(a, 1))
     dom_sects = sectors(axes(a, 2))
     (issorted(cod_sects) && allunique(cod_sects)) ||
@@ -301,7 +301,7 @@ end
 #     error — they've passed target axes incompatible with `m`).
 function delete_missing_sectors end
 
-function check_input(
+function TensorAlgebra.check_input(
         ::typeof(delete_missing_sectors),
         m::AbstractGradedMatrix, cod_ax::GradedOneTo, dom_ax::GradedOneTo
     )

src/tensoralgebra.jl

@@ -151,26 +151,32 @@ function tensor_product(r::SectorOneTo, s::TKS.Sector)
     return tensor_product(to_gradedrange(r), to_gradedrange(s))
 end
 
-# ========================  permutedimsopadd!  ========================
+# ========================  bipermutedimsopadd!  ========================
+# Primary overloads. The flat-perm permutedimsopadd! overloads forward here.
 
-function TensorAlgebra.permutedimsopadd!(
-        y::AbstractSectorArray, op, x::AbstractSectorArray, perm,
+function TensorAlgebra.bipermutedimsopadd!(
+        y::AbstractSectorArray, op, x::AbstractSectorArray,
+        perm_codomain, perm_domain,
         α::Number, β::Number
     )
-    sector(y) == permutedims(sector(x), perm) || throw(DimensionMismatch())
-    phase = fermion_permutation_phase(sector(x), perm)
-    TensorAlgebra.permutedimsopadd!(data(y), op, data(x), perm, phase * α, β)
+    check_input(bipermutedimsopadd!, y, x, perm_codomain, perm_domain)
+    phase = fermion_permutation_phase(sector(x), (perm_codomain..., perm_domain...))
+    bipermutedimsopadd!(
+        data(y), op, data(x), perm_codomain, perm_domain, phase * α, β
+    )
     return y
 end
 
-function TensorAlgebra.permutedimsopadd!(
-        y::AbstractGradedArray{<:Any, N}, op, x::AbstractGradedArray{<:Any, N}, perm,
+function TensorAlgebra.bipermutedimsopadd!(
+        y::AbstractGradedArray{<:Any, N}, op, x::AbstractGradedArray{<:Any, N},
+        perm_codomain, perm_domain,
         α::Number, β::Number
     ) where {N}
+    check_input(bipermutedimsopadd!, y, x, perm_codomain, perm_domain)
     # `scale!(y, 0)` doesn't reliably zero `y`: if any block of `y` holds
     # `NaN`/`Inf` (uninitialized memory from `undef` allocation or a stale
     # garbage value), `NaN * 0 == NaN` keeps it poisoned, and subsequent
-    # `permutedimsopadd!(..., α, one(α))` calls on a block of `y` that
+    # `bipermutedimsopadd!(..., α, one(β))` calls on a block of `y` that
     # doesn't get visited by the loop below would leak that garbage into the
     # result. Allocating broadcasts like `3 * a` go through this path (they
     # call with β == 0 on a fresh `similar`-allocated array); before this
@@ -179,10 +185,10 @@ function TensorAlgebra.permutedimsopadd!(
     iszero(β) ? FI.zero!(y) : scale!(y, β)
     for bI in eachblockstoredindex(x)
         b = Tuple(bI)
-        b_dest = Block(ntuple(i -> b[perm[i]], N))
+        b_dest = Block(ntuple(i -> b[(perm_codomain..., perm_domain...)[i]], N))
         y_b = view(y, Tuple(b_dest)...)
         x_b = x[bI]
-        TensorAlgebra.permutedimsopadd!(y_b, op, x_b, perm, α, one(α))
+        bipermutedimsopadd!(y_b, op, x_b, perm_codomain, perm_domain, α, one(β))
     end
     return y
 end

test/Project.toml

@@ -36,7 +36,7 @@ SUNRepresentations = "0.3"
 SafeTestsets = "0.1"
 SparseArraysBase = "0.9"
 Suppressor = "0.2.8"
-TensorAlgebra = "0.8"
+TensorAlgebra = "0.8, 0.9"
 TensorKitSectors = "0.3"
 Test = "1.10"
 TestExtras = "0.3.1"