diff --git a/Project.toml b/Project.toml index df7e71b1..bc6c8fbf 100644 --- a/Project.toml +++ b/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 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" diff --git a/docs/Project.toml b/docs/Project.toml index a1c794fd..6bb4067a 100644 --- a/docs/Project.toml +++ b/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" diff --git a/examples/Project.toml b/examples/Project.toml index 61c8cda5..96de614b 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -5,4 +5,4 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" path = ".." [compat] -TensorAlgebra = "0.13" +TensorAlgebra = "0.14" diff --git a/ext/TensorAlgebraGPUArraysCoreExt/TensorAlgebraGPUArraysCoreExt.jl b/ext/TensorAlgebraGPUArraysCoreExt/TensorAlgebraGPUArraysCoreExt.jl deleted file mode 100644 index 0af91528..00000000 --- a/ext/TensorAlgebraGPUArraysCoreExt/TensorAlgebraGPUArraysCoreExt.jl +++ /dev/null @@ -1,8 +0,0 @@ -module TensorAlgebraGPUArraysCoreExt - -import TensorAlgebra as TA -using GPUArraysCore: AnyGPUArray - -TA.iscpu(::AnyGPUArray) = false - -end diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 533519a2..55b3fdbf 100644 --- a/src/TensorAlgebra.jl +++ b/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") diff --git a/src/diagonal.jl b/src/diagonal.jl new file mode 100644 index 00000000..f836d931 --- /dev/null +++ b/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 diff --git a/src/permutedimsadd.jl b/src/permutedimsadd.jl index 04372c7a..70d63520 100644 --- a/src/permutedimsadd.jl +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index a46dc4da..527c0bc8 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/test_diagonal.jl b/test/test_diagonal.jl new file mode 100644 index 00000000..966e67ec --- /dev/null +++ b/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