Add bipermutedimsopadd! overloads for GradedArrays (companion to ITensor/TensorAlgebra.jl#158)

Author: mtfishman

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]

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!,
+    matricize, matricize_axes, permutedimsadd!, permutedimsopadd!, tensor_product_axis,
+    trivial_axis, trivialbiperm, tryflattenlinear, unmatricize
 using TensorKitSectors: TensorKitSectors as TKS
 using TypeParameterAccessors: type_parameters, unspecify_type_parameters
 

src/tensoralgebra.jl

@@ -151,26 +151,33 @@ 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
     )
+    perm = (perm_codomain..., perm_domain...)
     sector(y) == permutedims(sector(x), perm) || throw(DimensionMismatch())
     phase = fermion_permutation_phase(sector(x), perm)
-    TensorAlgebra.permutedimsopadd!(data(y), op, data(x), perm, phase * α, β)
+    TensorAlgebra.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}
+    perm = (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
@@ -182,7 +189,32 @@ function TensorAlgebra.permutedimsopadd!(
         b_dest = Block(ntuple(i -> b[perm[i]], N))
         y_b = view(y, Tuple(b_dest)...)
         x_b = x[bI]
-        TensorAlgebra.permutedimsopadd!(y_b, op, x_b, perm, α, one(α))
+        TensorAlgebra.bipermutedimsopadd!(
+            y_b,
+            op,
+            x_b,
+            perm_codomain,
+            perm_domain,
+            α,
+            one(α)
+        )
     end
     return y
 end
+
+# ========================  permutedimsopadd!  ========================
+# Flat-perm overloads forward to bipermutedimsopadd! with perm_domain = ().
+
+function TensorAlgebra.permutedimsopadd!(
+        y::AbstractSectorArray, op, x::AbstractSectorArray, perm,
+        α::Number, β::Number
+    )
+    return TensorAlgebra.bipermutedimsopadd!(y, op, x, perm, (), α, β)
+end
+
+function TensorAlgebra.permutedimsopadd!(
+        y::AbstractGradedArray{<:Any, N}, op, x::AbstractGradedArray{<:Any, N}, perm,
+        α::Number, β::Number
+    ) where {N}
+    return TensorAlgebra.bipermutedimsopadd!(y, op, x, perm, (), α, β)
+end