diff --git a/Project.toml b/Project.toml index 76bfe1ab..e1cbdc72 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.23.3" +version = "0.23.4" authors = ["ITensor developers and contributors"] [workspace] diff --git a/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl b/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl index 0dd50baa..3f493bc9 100644 --- a/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl +++ b/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl @@ -31,20 +31,16 @@ function TA.contractpermopadd!( α::Number, β::Number ) TA.check_input( - TA.contract!, + TA.contractpermopadd!, a_dest, perm_dest_codomain, perm_dest_domain, - a1, perm1_codomain, perm1_domain, - a2, perm2_codomain, perm2_domain + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain ) permblocks1 = Tuple.((perm1_codomain, perm1_domain)) permblocks2 = Tuple.((perm2_codomain, perm2_domain)) permblocks_dest = Tuple.((perm_dest_codomain, perm_dest_domain)) - conj1 = op1 === conj - conj2 = op2 === conj - a1′ = (op1 === identity || op1 === conj) ? a1 : op1.(a1) - a2′ = (op2 === identity || op2 === conj) ? a2 : op2.(a2) return TO.tensorcontract!( - a_dest, a1′, permblocks1, conj1, a2′, permblocks2, conj2, + a_dest, a1, permblocks1, op1 === conj, a2, permblocks2, op2 === conj, permblocks_dest, α, β, backend(algorithm), allocator(algorithm) ) end diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index d5f03e6b..f04fb258 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -5,7 +5,7 @@ export contract, contract!, contractalign, dual, isdual, MatrixAlgebra if VERSION >= v"1.11.0-DEV.469" eval( Meta.parse( - "public AbstractAlgorithm, add!, AddBroadcasted, addends, allocate_output, allocate_project, arguments, axes, bipartition, bipartition_axes, biperm, bipermutedims, bipermutedims!, bipermutedimsopadd!, cat_axis, cat_similar, check_input, concatenate, concatenate!, ConjBroadcasted, contractadd!, ContractAlgorithm, contractopadd!, contractperm, contractperm!, contractpermadd!, contractpermalign, contractpermopadd!, data, datatype, default_algorithm, dims2cat, directsum, dot, dotperm, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc, eigh_vals, fill_map, flattenlinear, has_bipartition, infer_aux_space, invsqrth_safe, is_output_view, is_projected, isidentitybiperm, label_type, left_null, left_orth, left_polar, LinearBroadcasted, linearbroadcasted, lq_compact, lq_full, matricize, MatricizeContract, matricizeop, matricizeop!, matricizeopcopy, matricizeopview, MATRIX_FUNCTIONS, ndims, ndims_codomain, ndims_domain, one, ones_map, operation, output_axes, PermutedDims, permuteddims, permutedims, permutedims!, permutedimsadd!, permutedimsop, permutedimsopadd!, project, project!, project_aux, project_hermitian, projectto!, qr_compact, qr_full, rand_map, randn_map, right_null, right_orth, right_polar, scalar, scale!, scaleadd!, ScaledBroadcasted, select_algorithm, similar_map, size, sqrth_invsqrth_safe, sqrth_safe, sum, svd_compact, svd_full, svd_trunc, svd_vals, TensorOperationsContract, to_range, tr, trivialrange, tryflattenlinear, tryproject, tryproject_aux, twist!, unchecked_project, unchecked_project_aux, ungrade, unmatricize, unmatricize!, unmatricize_factors, unmatricizeadd!, unproject, unscaled, zero!, zeros_map" + "public AbstractAlgorithm, add!, AddBroadcasted, addends, allocate_output, allocate_project, arguments, axes, bipartition, bipartition_axes, biperm, bipermutedims, bipermutedims!, bipermutedimsopadd!, cat_axis, cat_similar, check_input, concatenate, concatenate!, ConjArray, ConjBroadcasted, conjed, contractadd!, ContractAlgorithm, contractopadd!, contractperm, contractperm!, contractpermadd!, contractpermalign, contractpermop, contractpermopadd!, data, datatype, default_algorithm, dims2cat, directsum, dot, dotperm, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc, eigh_vals, fill_map, flattenlinear, has_bipartition, infer_aux_space, invsqrth_safe, is_output_view, is_projected, isidentitybiperm, label_type, left_null, left_orth, left_polar, LinearBroadcasted, linearbroadcasted, lq_compact, lq_full, matricize, MatricizeContract, matricizeop, matricizeop!, matricizeopcopy, matricizeopview, MATRIX_FUNCTIONS, ndims, ndims_codomain, ndims_domain, one, ones_map, operation, output_axes, PermutedDims, permuteddims, permutedims, permutedims!, permutedimsadd!, permutedimsop, permutedimsopadd!, project, project!, project_aux, project_hermitian, projectto!, qr_compact, qr_full, rand_map, randn_map, right_null, right_orth, right_polar, scalar, scale!, scaleadd!, ScaledBroadcasted, select_algorithm, similar_map, size, sqrth_invsqrth_safe, sqrth_safe, sum, svd_compact, svd_full, svd_trunc, svd_vals, TensorOperationsContract, to_range, tr, trivialrange, tryflattenlinear, tryproject, tryproject_aux, twist!, unchecked_project, unchecked_project_aux, ungrade, unmatricize, unmatricize!, unmatricize_factors, unmatricizeadd!, unproject, unscaled, zero!, zeros_map" ) ) end @@ -35,5 +35,6 @@ include("matrixfunctions.jl") include("similar_map.jl") include("projectto.jl") include("linearbroadcasted.jl") +include("conjarray.jl") end diff --git a/src/conjarray.jl b/src/conjarray.jl new file mode 100644 index 00000000..b9ec1d09 --- /dev/null +++ b/src/conjarray.jl @@ -0,0 +1,156 @@ +""" + ConjArray(a::AbstractArray) + +Lazy conjugate of `a`: an `AbstractArray` whose axes are the conjugated (dualized) parent axes +and whose elements are the conjugated parent elements. Constructed by [`conjed`](@ref). +Permutation, matricization and contraction absorb it by unwrapping to the parent and folding +`conj` into their `op`, so it is never materialized on its way into those operations. Whole-array +reads, writes and broadcasting go through the parent as well, as do reductions for a dense parent. +Scalar indexing is disallowed, so an operation without such a path errors instead of silently +looping over elements. +""" +struct ConjArray{T, N, P <: AbstractArray{T, N}} <: AbstractArray{T, N} + parent::P +end + +""" + conjed(a::AbstractArray) + +Lazy conjugate of `a`, returning a [`ConjArray`](@ref). The lazy companion to the eager +`conj(a)`, mirroring `adjoint`/`Adjoint`. `conjed(conjed(a))` returns `a`. +""" +conjed(a::AbstractArray) = ConjArray(a) +conjed(a::ConjArray) = parent(a) + +# Most methods below either go through an op-aware kernel (`permutedimsopadd!`, `matricizeop`, +# contraction), which owns what `conj` means for its array type, or use identities that hold for any +# antiunitary involution (`norm`, `iszero`, `dot` of two lazy conjugates). `fill!`, `==`, `mapreduce` +# and `dot` with one lazy operand instead assume `conj` acts elementwise on a flat memory layout, so +# they are only defined for a dense parent, and other array types opt in. +const ConjDenseArray{T, N} = ConjArray{T, N, <:DenseArray{T, N}} + +Base.parent(a::ConjArray) = a.parent + +# Conjugating an axis is identity for plain integer ranges and dualizes graded axes (where +# `conj` is overloaded as `dual` downstream). +Base.axes(a::ConjArray) = map(conj, axes(parent(a))) +Base.size(a::ConjArray) = size(parent(a)) + +function Base.getindex(a::ConjArray, I::Int...) + return throw( + ArgumentError( + "scalar indexing of a lazy `ConjArray` is disallowed, materialize it with `copy` first" + ) + ) +end +function Base.setindex!(a::ConjArray, value, I::Int...) + return throw( + ArgumentError( + "scalar indexing of a lazy `ConjArray` is disallowed, write the whole array with `copyto!`" + ) + ) +end + +# Allocate like the parent, so a graded parent yields a graded array. +function Base.similar(a::ConjArray, ::Type{T}, dims::Dims) where {T} + return similar(parent(a), T, dims) +end + +# Materialize through the eager `conj.(a)` path, so a lazy and an eager conjugate agree +# (including any parent-specific behavior such as a graded fermion sign). +Base.copy(a::ConjArray) = copy(ConjBroadcasted(parent(a))) +Base.Array{T, N}(a::ConjArray{<:Any, N}) where {T, N} = convert(Array{T, N}, copy(a)) +# Base `collect` iterates elements, which scalar indexing disallows. +Base.collect(a::ConjArray) = Array(a) +# Like `conj(::Adjoint)` in LinearAlgebra, this returns the parent rather than a copy. +Base.conj(a::ConjArray) = parent(a) + +function Base.copyto!(dest::AbstractArray, src::ConjArray) + return permutedimsopadd!( + dest, + conj, + parent(src), + ntuple(identity, ndims(src)), + true, + false + ) +end +# Write into the parent the conjugate of what the lazy conjugate should hold. `conj` is an +# involution, fermion signs included, so this is exact. +function Base.copyto!(dest::ConjArray, src::AbstractArray) + permutedimsopadd!(parent(dest), conj, src, ntuple(identity, ndims(src)), true, false) + return dest +end +# Also resolves the ambiguity between the two methods above. +function Base.copyto!(dest::ConjArray, src::ConjArray) + copyto!(parent(dest), parent(src)) + return dest +end +function Base.fill!(a::ConjDenseArray, value) + fill!(parent(a), conj(value)) + return a +end + +Base.:(==)(a::ConjArray, b::ConjArray) = parent(a) == parent(b) +# A fused two-array reduction rather than a copy, so it maps to a single kernel on GPU. +function Base.:(==)(a::ConjDenseArray, b::DenseArray) + axes(a) == axes(b) || return false + return mapreduce((x, y) -> conj(x) == y, &, parent(a), b; init = true) +end +Base.:(==)(a::DenseArray, b::ConjDenseArray) = b == a +# Conjugation, and any sign it carries, preserves every p-norm. `isapprox` goes through `norm`. +LinearAlgebra.norm(a::ConjArray, p::Real = 2) = LinearAlgebra.norm(parent(a), p) +Base.iszero(a::ConjArray) = iszero(parent(a)) +# `conj` is antiunitary, fermion signs included, so a pair of lazy conjugates reduces to the parents. +function LinearAlgebra.dot(a::ConjArray, b::ConjArray) + return conj(LinearAlgebra.dot(parent(a), parent(b))) +end +# Against a dense array, the lazy conjugate cancels the conjugation in `dot`, leaving the bilinear +# pairing of the parent with `b` (BLAS `dotu`). +LinearAlgebra.dot(a::ConjDenseArray, b::DenseArray) = transpose(vec(parent(a))) * vec(b) +LinearAlgebra.dot(a::DenseArray, b::ConjDenseArray) = conj(LinearAlgebra.dot(b, a)) +function Base.mapreduce(f, op, a::ConjDenseArray; kwargs...) + return mapreduce(f ∘ conj, op, parent(a); kwargs...) +end +Base.permutedims(a::ConjArray, perm) = permutedims(a, perm) + +Base.show(io::IO, a::ConjArray) = print(io, "conjed(", parent(a), ")") +function Base.show(io::IO, mime::MIME"text/plain", a::ConjArray) + print(io, "conjed of ") + return show(io, mime, copy(a)) +end + +# A broadcast over a lazy conjugate reads it as `conj.(parent)`, so it never indexes the wrapper +# and a linear one lowers to a `ConjBroadcasted` leaf. +Base.Broadcast.broadcastable(a::ConjArray) = Base.Broadcast.broadcasted(conj, parent(a)) +function Base.Broadcast.BroadcastStyle(::Type{<:ConjArray{<:Any, <:Any, P}}) where {P} + return Base.Broadcast.BroadcastStyle(P) +end + +linearbroadcasted(::typeof(conj), a::ConjArray) = parent(a) + +# Each operation unwraps a `ConjArray` to its parent and folds `conj` into `op` +# (`identity -> conj`, `conj -> identity`), before any backend dispatch. +function bipermutedimsopadd!( + dest, op, src::ConjArray, + perm_codomain, perm_domain, + α::Number, β::Number + ) + return bipermutedimsopadd!( + dest, _compose_op(op, conj), parent(src), + perm_codomain, perm_domain, + α, β + ) +end +function permutedimsop(op, src::ConjArray, perm_codomain, perm_domain) + return permutedimsop(_compose_op(op, conj), parent(src), perm_codomain, perm_domain) +end +function matricizeop(op, a::ConjArray, perm_codomain, perm_domain) + return matricizeop(_compose_op(op, conj), parent(a), perm_codomain, perm_domain) +end +function matricizeopcopy(op, a::ConjArray, perm_codomain, perm_domain) + return matricizeopcopy(_compose_op(op, conj), parent(a), perm_codomain, perm_domain) +end + +_unwrap_conj(op, a) = (op, a) +_unwrap_conj(op, a::ConjArray) = (_compose_op(op, conj), parent(a)) diff --git a/src/contract/allocate_output.jl b/src/contract/allocate_output.jl index 5fb32103..a504a8d1 100644 --- a/src/contract/allocate_output.jl +++ b/src/contract/allocate_output.jl @@ -6,33 +6,39 @@ function check_biperm(a, perm_codomain, perm_domain) return nothing end -function check_input( - ::typeof(contract), - a1, perm1_codomain, perm1_domain, - a2, perm2_codomain, perm2_domain - ) - check_biperm(a1, perm1_codomain, perm1_domain) - check_biperm(a2, perm2_codomain, perm2_domain) - # The contracted axes of `a1` (its domain) pair with the contracted axes of `a2` (its codomain), - # and each pair must be a dual pair: `dual(ax1) == ax2`. `dual` falls back to the identity for - # plain ranges, so this reduces to `ax1 == ax2` for non-graded arrays. +# The contracted axes of `op1(a1)` (its domain) pair with the contracted axes of `op2(a2)` (its +# codomain), and each pair must be a dual pair: `dual(ax1) == ax2`. `op` is `identity` or `conj`, +# and `conj` dualizes a graded axis. `dual` and `conj` fall back to the identity for plain ranges, +# so this reduces to `ax1 == ax2` for non-graded arrays. +function check_contracted_axes(op1, a1, perm1_domain, op2, a2, perm2_codomain) length(perm1_domain) == length(perm2_codomain) || throw( ArgumentError( "Number of contracted axes do not match: `a1` has $(length(perm1_domain)), `a2` has $(length(perm2_codomain))" ) ) for (i, j) in zip(perm1_domain, perm2_codomain) - ax1 = axes(a1, i) - ax2 = axes(a2, j) + ax1 = op1(axes(a1, i)) + ax2 = op2(axes(a2, j)) dual(ax1) == ax2 || throw( ArgumentError( - "Contracted axes do not match: `axes(a1, $i) = $ax1` and `axes(a2, $j) = $ax2`" + "Contracted axes do not match: `$op1(axes(a1, $i)) = $ax1` and `$op2(axes(a2, $j)) = $ax2`" ) ) end return nothing end +function check_input( + ::typeof(contract), + a1, perm1_codomain, perm1_domain, + a2, perm2_codomain, perm2_domain + ) + check_biperm(a1, perm1_codomain, perm1_domain) + check_biperm(a2, perm2_codomain, perm2_domain) + check_contracted_axes(identity, a1, perm1_domain, identity, a2, perm2_codomain) + return nothing +end + function check_input( ::typeof(contract!), a_dest, perm_dest_codomain, perm_dest_domain, @@ -53,16 +59,62 @@ function check_input( return nothing end -# TODO: Use `ArrayLayouts`-like `MulAdd` object, -# i.e. `ContractAdd`? +function check_input( + ::typeof(contractpermop), + perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain + ) + for op in (op1, op2) + op === identity || op === conj || + throw(ArgumentError("`op` must be `identity` or `conj`, got `$op`")) + end + check_biperm(a1, perm1_codomain, perm1_domain) + check_biperm(a2, perm2_codomain, perm2_domain) + check_contracted_axes(op1, a1, perm1_domain, op2, a2, perm2_codomain) + return nothing +end + +function check_input( + ::typeof(contractpermopadd!), + a_dest, perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain + ) + # TODO: FIXME: Check that uncontracted axes match. + check_input( + contractpermop, + perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain + ) + check_biperm(a_dest, perm_dest_codomain, perm_dest_domain) + return nothing +end + function output_axes( ::typeof(contract), perm_dest_codomain, perm_dest_domain, a1, perm1_codomain, perm1_domain, a2, perm2_codomain, perm2_domain ) - axes_codomain, _ = bipartition(axes(a1), perm1_codomain, perm1_domain) - _, axes_domain = bipartition(axes(a2), perm2_codomain, perm2_domain) + return output_axes( + contractpermop, + perm_dest_codomain, perm_dest_domain, + identity, a1, perm1_codomain, perm1_domain, + identity, a2, perm2_codomain, perm2_domain + ) +end +# TODO: Use `ArrayLayouts`-like `MulAdd` object, +# i.e. `ContractAdd`? +function output_axes( + ::typeof(contractpermop), + perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain + ) + axes_codomain = map(i -> op1(axes(a1, i)), perm1_codomain) + axes_domain = map(i -> op2(axes(a2, i)), perm2_domain) axes_uncontracted = (axes_codomain..., axes_domain...) axes_codomain_dest, axes_domain_dest = bipartition( axes_uncontracted, perm_dest_codomain, perm_dest_domain @@ -72,32 +124,42 @@ function output_axes( return axes_codomain_dest, conj.(axes_domain_dest) end -# TODO: Use `ArrayLayouts`-like `MulAdd` object, -# i.e. `ContractAdd`? -# The destination `contract` writes into. A structured operand type overloads this directly, deriving -# the axes and element type from `output_axes` and `Base.promote_op` as below; the permutations are -# part of the signature because the contraction pattern is not recoverable from the destination leg -# counts alone. function allocate_output( ::typeof(contract), perm_dest_codomain, perm_dest_domain, a1, perm1_codomain, perm1_domain, a2, perm2_codomain, perm2_domain ) + return allocate_output( + contractpermop, + perm_dest_codomain, perm_dest_domain, + identity, a1, perm1_codomain, perm1_domain, + identity, a2, perm2_codomain, perm2_domain + ) +end +# TODO: Use `ArrayLayouts`-like `MulAdd` object, +# i.e. `ContractAdd`? +# The destination `contractpermop` writes into. A structured operand type overloads this directly, +# deriving the axes and element type from `output_axes` and `Base.promote_op` as below; the +# permutations are part of the signature because the contraction pattern is not recoverable from the +# destination leg counts alone. +function allocate_output( + ::typeof(contractpermop), + perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain + ) check_input( - contract, - a1, - perm1_codomain, - perm1_domain, - a2, - perm2_codomain, - perm2_domain + contractpermop, + perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain ) axes_codomain_dest, axes_domain_dest = output_axes( - contract, + contractpermop, perm_dest_codomain, perm_dest_domain, - a1, perm1_codomain, perm1_domain, - a2, perm2_codomain, perm2_domain + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain ) T = Base.promote_op(matprod, eltype(a1), eltype(a2)) return zero!(similar_map(a1, T, axes_codomain_dest, axes_domain_dest)) diff --git a/src/contract/contract.jl b/src/contract/contract.jl index 7483251b..88fff2eb 100644 --- a/src/contract/contract.jl +++ b/src/contract/contract.jl @@ -107,17 +107,34 @@ function contractpermalign( a2, perm2_codomain, perm2_domain; kwargs... ) + return contractpermop( + perm_dest_codomain, perm_dest_domain, + identity, a1, perm1_codomain, perm1_domain, + identity, a2, perm2_codomain, perm2_domain; + kwargs... + ) +end +function contractpermop( + perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain; + kwargs... + ) + # Unwrap lazy conjugates before allocating, so the destination is allocated for the operands' + # own array types. + op1, a1 = _unwrap_conj(op1, a1) + op2, a2 = _unwrap_conj(op2, a2) a_dest = allocate_output( - contract, + contractpermop, perm_dest_codomain, perm_dest_domain, - a1, perm1_codomain, perm1_domain, - a2, perm2_codomain, perm2_domain + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain ) - return contractperm!( + return contractpermopadd!( a_dest, perm_dest_codomain, perm_dest_domain, - a1, perm1_codomain, perm1_domain, - a2, perm2_codomain, perm2_domain; - kwargs... + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain, + true, false; kwargs... ) end @@ -212,7 +229,10 @@ function contractpermopadd!( alg = nothing ) # Input validation belongs to the algorithm's kernel below, which is a public rung callable - # directly, so this entry only selects and dispatches. + # directly, so this entry only unwraps lazy conjugates, selects and dispatches. Unwrapping + # before selection lets a backend kernel see its own array type with `op = conj`. + op1, a1 = _unwrap_conj(op1, a1) + op2, a2 = _unwrap_conj(op2, a2) algorithm = select_algorithm(contract!, alg, (a_dest, a1, a2)) return contractpermopadd!( algorithm, diff --git a/src/contract/contract_matricize.jl b/src/contract/contract_matricize.jl index 67973a38..f52132a4 100644 --- a/src/contract/contract_matricize.jl +++ b/src/contract/contract_matricize.jl @@ -15,10 +15,10 @@ function contractpermopadd!( invperm_codomain, invperm_domain = bipartition(invperm(biperm_dest), Val(length(biperm1_codomain))) check_input( - contract!, + contractpermopadd!, a_dest, invperm_codomain, invperm_domain, - a1, biperm1_codomain, biperm1_domain, - a2, biperm2_codomain, biperm2_domain + op1, a1, biperm1_codomain, biperm1_domain, + op2, a2, biperm2_codomain, biperm2_domain ) a1_mat = matricizeop(op1, a1, biperm1_codomain, biperm1_domain) a2_mat = matricizeop(op2, a2, biperm2_codomain, biperm2_domain) diff --git a/src/diagonal.jl b/src/diagonal.jl index 3ce43938..64d1f8e7 100644 --- a/src/diagonal.jl +++ b/src/diagonal.jl @@ -80,19 +80,22 @@ end # not match this signature, falling back to the generic dense allocation the way `Diagonal`/dense # mixing does. function allocate_output( - ::typeof(contract), + ::typeof(contractpermop), perm_dest_codomain::Tuple{Int}, perm_dest_domain::Tuple{Int}, - a1::Diagonal, perm1_codomain, perm1_domain, - a2::Diagonal, perm2_codomain, perm2_domain + op1, a1::Diagonal, perm1_codomain, perm1_domain, + op2, a2::Diagonal, perm2_codomain, perm2_domain ) check_input( - contract, a1, perm1_codomain, perm1_domain, a2, perm2_codomain, perm2_domain + contractpermop, + perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain ) axes_codomain_dest, _ = output_axes( - contract, + contractpermop, perm_dest_codomain, perm_dest_domain, - a1, perm1_codomain, perm1_domain, - a2, perm2_codomain, perm2_domain + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain ) T = Base.promote_op(matprod, eltype(a1), eltype(a2)) return Diagonal(zero!(similar(a1.diag, T, (only(axes_codomain_dest),)))) diff --git a/test/test_conjarray.jl b/test/test_conjarray.jl new file mode 100644 index 00000000..e1c016c8 --- /dev/null +++ b/test/test_conjarray.jl @@ -0,0 +1,175 @@ +using LinearAlgebra: Diagonal, dot, norm +using TensorAlgebra: TensorAlgebra as TA, ConjArray, bipermutedimsopadd!, conjed, contract, + contractopadd!, matricize, permutedims, qr_compact +using Test: @test, @test_throws, @testset + +@testset "ConjArray (plain arrays)" begin + a = randn(ComplexF64, 3, 4) + c = conjed(a) + @test c isa ConjArray + @test parent(c) === a + @test conjed(c) === a + @test eltype(c) === eltype(a) + @test size(c) == size(a) + @test axes(c) == axes(a) + @test copy(c) == conj(a) + @test conj(c) === a +end + +@testset "ConjArray goes through its parent instead of indexing" begin + a = randn(ComplexF64, 3, 4) + b = randn(ComplexF64, 3, 4) + c = conjed(a) + @test_throws ArgumentError c[2, 3] + @test_throws ArgumentError c[2, 3] = 1 + @test collect(c) == conj(a) + + @test c == conj(a) + @test conj(a) == c + @test c == conjed(copy(a)) + @test c ≈ conj(a) + @test norm(c) == norm(a) + @test !iszero(c) + @test iszero(conjed(zero(a))) + @test dot(c, b) ≈ dot(conj(a), b) + @test dot(b, c) ≈ dot(b, conj(a)) + @test dot(c, conjed(b)) ≈ dot(conj(a), conj(b)) + d = Diagonal(randn(ComplexF64, 3)) + m = randn(ComplexF64, 3, 3) + @test_throws ArgumentError dot(conjed(m), d) + @test_throws ArgumentError dot(d, conjed(m)) + # Elementwise methods are only defined for a dense parent. + @test_throws ArgumentError dot(conjed(d), m) + @test_throws ArgumentError conjed(d) == conj(d) + @test_throws ArgumentError sum(conjed(d)) + @test_throws ArgumentError fill!(conjed(copy(d)), 1) + @test Array(c) == conj(a) + @test sum(c) ≈ sum(conj(a)) + @test maximum(real, c) == maximum(real, a) + @test Base.permutedims(c, (2, 1)) == permutedims(conj(a), (2, 1)) + @test c .+ b == conj(a) .+ b + @test abs.(c) == abs.(a) + @test sprint(show, c) == "conjed($(sprint(show, a)))" + @test startswith(sprint(show, MIME"text/plain"(), c), "conjed of ") + + # Writes go to the parent as the conjugate of what is written. + d = conjed(zeros(ComplexF64, 3, 4)) + @test copyto!(d, b) === d + @test parent(d) == conj(b) + @test copyto!(similar(b), d) == b + @test fill!(d, 2im) === d + @test parent(d) == fill(-2im, 3, 4) + @test TA.zero!(d) === d + @test iszero(parent(d)) + @test similar(c, Float64, (2, 5)) isa Matrix{Float64} +end + +@testset "factorizations of a ConjArray" begin + a = randn(ComplexF64, 2, 3, 4) + Q, R = qr_compact(conjed(a), (1, 2), (3,)) + Qe, Re = qr_compact(conj(a), (1, 2), (3,)) + @test Q ≈ Qe + @test R ≈ Re +end + +@testset "ConjArray is absorbed by permutation and matricize" begin + a = randn(ComplexF64, 2, 3, 4, 5) + c = conjed(a) + @test permutedims(c, (3, 1, 4, 2)) == permutedims(conj(a), (3, 1, 4, 2)) + for (pc, pd) in (((1, 2, 3, 4), ()), ((2, 4), (1, 3))) + dest = zeros(ComplexF64, size(a)[[pc..., pd...]]...) + bipermutedimsopadd!(dest, identity, c, pc, pd, true, false) + @test dest == permutedims(conj(a), (pc..., pd...)) + dest = zeros(ComplexF64, size(a)[[pc..., pd...]]...) + bipermutedimsopadd!(dest, conj, c, pc, pd, true, false) + @test dest == permutedims(a, (pc..., pd...)) + end + @test matricize(c, (2, 4), (1, 3)) == matricize(conj(a), (2, 4), (1, 3)) + @test matricize(c, Val(2)) == matricize(conj(a), Val(2)) +end + +@testset "ConjArray is absorbed by contraction" begin + a = randn(ComplexF64, 2, 3, 4) + b = randn(ComplexF64, 4, 3, 5) + labels_a, labels_b = (:i, :j, :k), (:k, :j, :l) + for (x, y) in ((conjed(a), b), (a, conjed(b)), (conjed(a), conjed(b))) + xe = x isa ConjArray ? copy(x) : x + ye = y isa ConjArray ? copy(y) : y + r, labels_r = contract(x, labels_a, y, labels_b) + re, labels_re = contract(xe, labels_a, ye, labels_b) + @test labels_r == labels_re + @test r == re + end + + # The conjugate is folded into the op passed alongside it. + dest = zeros(ComplexF64, 2, 5) + contractopadd!( + dest, + (:i, :l), + conj, + conjed(a), + labels_a, + identity, + b, + labels_b, + true, + false + ) + @test dest == first(contract(a, labels_a, b, labels_b)) +end + +@testset "contractpermop allocates for the unwrapped operands" begin + a = randn(ComplexF64, 2, 3) + b = randn(ComplexF64, 3, 4) + @test TA.contractpermop((1,), (2,), conj, a, (1,), (2,), identity, b, (1,), (2,)) ≈ + conj(a) * b + @test TA.contractpermop( + (1,), + (2,), + identity, + conjed(a), + (1,), + (2,), + identity, + b, + (1,), + (2,) + ) ≈ + conj(a) * b + + # Allocation sees the parent, so a `Diagonal` pair still allocates a `Diagonal`. + d1 = Diagonal(randn(ComplexF64, 3)) + d2 = Diagonal(randn(ComplexF64, 3)) + r, _ = contract(conjed(d1), (:i, :j), d2, (:j, :k)) + @test r isa Diagonal + @test r ≈ conj(d1) * d2 +end + +@testset "op-aware contraction check_input" begin + a = randn(ComplexF64, 2, 3) + b = randn(ComplexF64, 3, 4) + @test isnothing( + TA.check_input( + TA.contractpermopadd!, + zeros(ComplexF64, 2, 4), + (1,), + (2,), + conj, + a, + (1,), + (2,), + identity, + b, + (1,), + (2,) + ) + ) + @test_throws ArgumentError TA.check_input( + TA.contractpermopadd!, zeros(ComplexF64, 2, 4), (1,), (2,), + abs, a, (1,), (2,), identity, b, (1,), (2,) + ) + @test_throws ArgumentError TA.check_input( + TA.contractpermopadd!, zeros(ComplexF64, 2, 4), (1,), (2,), + conj, a, (1,), (2,), identity, randn(4, 4), (1,), (2,) + ) +end diff --git a/test/test_exports.jl b/test/test_exports.jl index 67e6b207..02b93b3f 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -21,10 +21,12 @@ using Test: @test, @testset :addends, :allocate_output, :allocate_project, :arguments, :axes, :bipartition, :bipartition_axes, :biperm, :bipermutedims, :bipermutedims!, :bipermutedimsopadd!, :cat_axis, :cat_similar, - :check_input, :concatenate, :concatenate!, :ConjBroadcasted, + :check_input, :concatenate, :concatenate!, :ConjArray, :ConjBroadcasted, + :conjed, :contractadd!, :ContractAlgorithm, :contractopadd!, :contractperm, :contractperm!, - :contractpermadd!, :contractpermalign, :contractpermopadd!, :data, + :contractpermadd!, :contractpermalign, :contractpermop, + :contractpermopadd!, :data, :datatype, :default_algorithm, :dims2cat, :directsum, :dot, :dotperm, :eig_full, :eig_trunc, diff --git a/test/test_tensorkitext.jl b/test/test_tensorkitext.jl index 10e3b060..9ae40bd3 100644 --- a/test/test_tensorkitext.jl +++ b/test/test_tensorkitext.jl @@ -26,6 +26,23 @@ using Test: @test, @test_throws, @testset @test c ≈ a * b end + @testset "contract with a conjugated operand" begin + W = Rep[U₁](0 => 2, 1 => 1) + X = Rep[U₁](0 => 1, 1 => 2) + Y = Rep[U₁](-1 => 1, 0 => 2) + a = randn(rng, elt, W, X) + b = randn(rng, elt, W, Y) + # `conj(a)` dualizes the spaces, so its first leg contracts with `b`'s codomain leg `W`. + ref, labels = contract( + copy(TensorAlgebra.ConjBroadcasted(a)), (:i, :j), b, (:i, :k) + ) + c = zero(ref) + TensorAlgebra.contractopadd!( + c, labels, conj, a, (:i, :j), identity, b, (:i, :k), true, false + ) + @test c ≈ ref + end + @testset "contract abelian, rank-3 with permuted output" begin A1 = Rep[U₁](0 => 2, 1 => 1) A2 = Rep[U₁](0 => 1, 1 => 1)