Simpler StridedView wrapping and Diagonal support

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.13.4"
+version = "0.14.0"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -15,23 +15,20 @@ StridedViews = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 
 [weakdeps]
-GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 
 [extensions]
-TensorAlgebraGPUArraysCoreExt = "GPUArraysCore"
 TensorAlgebraMooncakeExt = "Mooncake"
 TensorAlgebraTensorOperationsExt = "TensorOperations"
 
 [compat]
 EllipsisNotation = "1.8"
-GPUArraysCore = "0.2"
 LinearAlgebra = "1.10"
 MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"
 Mooncake = "0.4.202, 0.5"
-Strided = "2.3.5"
-StridedViews = "0.4.1, 0.5"
+Strided = "2.6"
+StridedViews = "0.5"
 TensorOperations = "5"
 TupleTools = "1.6"
 julia = "1.10"

docs/Project.toml

@@ -11,4 +11,4 @@ path = ".."
 Documenter = "1.8.1"
 ITensorFormatter = "0.2.27"
 Literate = "2.20.1"
-TensorAlgebra = "0.13"
+TensorAlgebra = "0.14"

examples/Project.toml

@@ -5,4 +5,4 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 path = ".."
 
 [compat]
-TensorAlgebra = "0.13"
+TensorAlgebra = "0.14"

src/TensorAlgebra.jl

@@ -19,6 +19,7 @@ include("bituple.jl")
 include("permutedimsadd.jl")
 include("conjarray.jl")
 include("matricize.jl")
+include("diagonal.jl")
 include("to_range.jl")
 include("contract/contractalgorithm.jl")
 include("contract/contract.jl")

src/diagonal.jl

@@ -0,0 +1,49 @@
+using LinearAlgebra: Diagonal
+
+# `Diagonal` participates in the `ReshapeFusion` interface like a dense matrix (it fuses with
+# the same row/column reshape order), but its structure is preserved wherever the result of
+# an operation is still diagonal. These methods hook the lowest-level primitives, so the
+# convenience wrappers built on them (`bipermutedims`, `permutedimsadd!`, `add!`, and the
+# matrix functions, which all route through `bipermutedimsopadd!` and `allocate_output`)
+# preserve `Diagonal`. Structure is given up (via the generic reshape path) only where the
+# result genuinely is not diagonal: vectorizing matricizations and `Diagonal`/dense mixing
+# (the latter falls back to Base's dense `similar`, since `contract` allocates from a flat
+# axis tuple, not a `BiTuple`).
+
+# Permuting the two axes of a square `Diagonal` (identity or transpose) leaves the stored
+# diagonal unchanged, so accumulate straight onto it.
+function bipermutedimsopadd!(
+        dest::Diagonal, op, src::Diagonal,
+        perm_codomain, perm_domain,
+        α::Number, β::Number
+    )
+    check_input(bipermutedimsopadd!, dest, op, src, perm_codomain, perm_domain)
+    _opadd!(dest.diag, op, src.diag, α, β)
+    return dest
+end
+
+# The bipermutation of a square `Diagonal` is again a square `Diagonal` of the same size (it
+# only swaps or keeps the two axes), so allocate the `permutedimsop`/`bipermutedims` output as
+# a `Diagonal`. The squareness comes from `src`, not from the axes: an axis-based
+# `similar(::BiTuple)` could not preserve it, since row/column axes alone do not encode that
+# the result is square.
+function allocate_output(
+        ::typeof(permutedimsop),
+        op,
+        src::Diagonal,
+        perm_codomain,
+        perm_domain
+    )
+    T = Base.promote_op(op, eltype(src))
+    return Diagonal(similar(src.diag, T))
+end
+
+# A `Diagonal` is already a matrix; the `(1 codomain, 1 domain)` matricization is the identity
+# reshape, so return it directly (maybe-alias, matching `matricize`'s general contract).
+matricize(::ReshapeFusion, a::Diagonal, ::Val{1}) = a
+function unmatricize(
+        ::ReshapeFusion, m::Diagonal,
+        ::Tuple{<:AbstractUnitRange}, ::Tuple{<:AbstractUnitRange}
+    )
+    return m
+end

src/permutedimsadd.jl

@@ -9,14 +9,6 @@ hook: downstream array types can overload it to return a custom lazy permuted-di
 """
 permuteddims(a::AbstractArray, perm) = PermutedDimsArray(a, perm)
 
-# Specify if an array is on CPU. This is helpful for backends that don't support
-# operations on GPU, such as Strided.jl.
-iscpu(::AbstractArray) = true
-# Convert to StridedView only if all arrays are strided and on CPU.
-function maybestrided(as::AbstractArray...)
-    return all(a -> SV.isstrided(a) && iscpu(a), as) ? SV.StridedView.(as) : as
-end
-
 # ---------------------------------------------------------------------------- #
 # bipermutedimsopadd! — the primary materialization primitive
 # ---------------------------------------------------------------------------- #
@@ -66,34 +58,24 @@ function bipermutedimsopadd!(
     perm = (perm_codomain..., perm_domain...)
     check_input(bipermutedimsopadd!, dest, op, src, perm_codomain, perm_domain)
 
-    # 0-dim short-circuit: avoid the permute-broadcast path entirely so that
-    # downstream array types (e.g. `BlockSparseArray{T, 0}`) don't have to define
-    # `getindex` on a 0-dim `PermutedDimsArray` wrapper around them.
-    # The `iszero(β)` guard follows the BLAS convention that `β = 0` means `dest`
-    # is write-only — its slot need not be defined. This matters for element types
-    # whose `undef` storage is unreadable, e.g. `Array{BigFloat, 0}(undef)[]` throws
-    # `UndefRefError`.
-    if iszero(ndims(dest))
-        if iszero(β)
-            dest[] = α * op(src[])
-        else
-            dest[] = β * dest[] + α * op(src[])
-        end
-        return dest
-    end
+    dest′ = SV.StridedView(dest)
+    src′ = permutedims(SV.StridedView(src), perm)
+    _opadd!(dest′, op, src′, α, β)
+    return dest
+end
 
-    dest′, src′ = maybestrided(dest, permuteddims(src, perm))
+function _opadd!(dest::AbstractArray, op, src::AbstractArray, α, β)
     if op === identity
         if iszero(β)
-            dest′ .= α .* src′
+            dest .= α .* src
         else
-            dest′ .= β .* dest′ .+ α .* src′
+            dest .= β .* dest .+ α .* src
         end
     else
         if iszero(β)
-            dest′ .= α .* op.(src′)
+            dest .= α .* op.(src)
         else
-            dest′ .= β .* dest′ .+ α .* op.(src′)
+            dest .= β .* dest .+ α .* op.(src)
         end
     end
     return dest

test/Project.toml

@@ -36,7 +36,7 @@ Random = "1.10"
 SafeTestsets = "0.1"
 StableRNGs = "1.0.2"
 Suppressor = "0.2"
-TensorAlgebra = "0.13"
+TensorAlgebra = "0.14"
 TensorOperations = "5.1.4"
 Test = "1.10"
 TestExtras = "0.3.1"

test/test_diagonal.jl

@@ -0,0 +1,78 @@
+using LinearAlgebra: Diagonal, diag
+using TensorAlgebra: TensorAlgebra
+using Test: @test, @testset
+
+@testset "Diagonal TensorAlgebra interface (eltype=$elt)" for elt in (Float64, ComplexF64)
+    d = Diagonal(elt[2, 3, 4])
+
+    @testset "bipermutedims preserves Diagonal" begin
+        b1 = TensorAlgebra.bipermutedims(d, (1,), (2,))
+        @test b1 isa Diagonal
+        @test b1 !== d
+        @test b1 == d
+        # A 2D transpose of a Diagonal is the same Diagonal.
+        b2 = TensorAlgebra.bipermutedims(d, (2,), (1,))
+        @test b2 isa Diagonal
+        @test b2 == d
+    end
+
+    @testset "permutedimsop applies the op to the data" begin
+        dz = Diagonal(elt <: Complex ? elt[1 + 2im, 3 - im, 2im] : elt[1, 3, 2])
+        p = TensorAlgebra.permutedimsop(conj, dz, (1,), (2,))
+        @test p isa Diagonal
+        @test p == conj(dz)
+    end
+
+    @testset "allocate_output returns a Diagonal of the same size" begin
+        out = TensorAlgebra.allocate_output(
+            TensorAlgebra.permutedimsop, identity, d, (1,), (2,)
+        )
+        @test out isa Diagonal
+        @test size(out) == size(d)
+        @test eltype(out) === elt
+    end
+
+    @testset "add! accumulates onto a Diagonal" begin
+        dest = Diagonal(elt[1, 1, 1])
+        # `add!(dest, src, α, β)` computes `α * src + β * dest`.
+        TensorAlgebra.add!(dest, d, elt(2), elt(1))
+        @test dest isa Diagonal
+        @test dest == Diagonal(elt[5, 7, 9])
+    end
+
+    @testset "matricize(1, 1) is the identity reshape" begin
+        m = TensorAlgebra.matricize(TensorAlgebra.ReshapeFusion(), d, Val(1))
+        @test m === d
+    end
+
+    @testset "unmatricize round-trips a Diagonal" begin
+        ax = axes(d, 1)
+        back = TensorAlgebra.unmatricize(TensorAlgebra.ReshapeFusion(), d, (ax,), (ax,))
+        @test back === d
+    end
+
+    @testset "matrix functions preserve Diagonal" begin
+        dp = Diagonal(elt[4, 9, 16])
+        s = TensorAlgebra.sqrt(dp, ("i", "j"), ("i",), ("j",))
+        @test s isa Diagonal
+        @test s ≈ sqrt(dp)
+        e = TensorAlgebra.exp(dp, ("i", "j"), ("i",), ("j",))
+        @test e isa Diagonal
+        @test e ≈ exp(dp)
+    end
+
+    @testset "contract densifies (Diagonal is an input structure, not an output one)" begin
+        d2 = Diagonal(elt[10, 20, 30])
+        # One contracted leg: a matrix product, materialized dense.
+        c2, = TensorAlgebra.contract(d, ("i", "k"), d2, ("k", "j"))
+        @test !(c2 isa Diagonal)
+        @test c2 ≈ d * d2
+        # Both legs contracted: a scalar.
+        c0, = TensorAlgebra.contract(d, ("i", "j"), d2, ("i", "j"))
+        @test ndims(c0) == 0
+        @test c0[] ≈ sum(diag(d) .* diag(d2))
+        # No contracted legs: a rank-4 outer product.
+        c4, = TensorAlgebra.contract(d, ("i", "j"), d2, ("k", "l"))
+        @test ndims(c4) == 4
+    end
+end