diff --git a/Project.toml b/Project.toml index 63a6acda..0f7f4e7c 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.19.2" +version = "0.19.3" authors = ["ITensor developers and contributors"] [workspace] diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 6b73a832..f386638a 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -315,11 +315,31 @@ function Base.copyto!(dest::AbstractTensorMap, src::TensorAlgebra.LinearBroadcas return TensorAlgebra.add!(dest, src, true, false) end -function Base.copy(::Base.Broadcast.Broadcasted{TensorMapStyle}) - return error( +# Allocation for a linear-combination `copy`/`copyto!`: seed the result off a `TensorMap` operand +# (`broadcast_prototype`) so it inherits that operand's storage type, and take the axes from the +# flattened expression, which carry the conj-dualization that Base's `combine_axes` drops. +function Base.similar(bc::Base.Broadcast.Broadcasted{TensorMapStyle}, ::Type{T}) where {T} + lb = TensorAlgebra.flattenlinear(bc) + return TensorAlgebra.similar_map(broadcast_prototype(lb), T, axes(lb), ()) +end +# The first `TensorMap` operand in a broadcast expression, unwrapping the linear-fold leaves. +broadcast_prototype(a::AbstractTensorMap) = a +broadcast_prototype(a::TensorAlgebra.PermutedDims) = broadcast_prototype(parent(a)) +function broadcast_prototype(a::TensorAlgebra.ScaledBroadcasted) + return broadcast_prototype(TensorAlgebra.unscaled(a)) +end +broadcast_prototype(a::TensorAlgebra.ConjBroadcasted) = broadcast_prototype(parent(a)) +function broadcast_prototype(a::TensorAlgebra.AddBroadcasted) + return broadcast_prototype(first(TensorAlgebra.addends(a))) +end + +function Base.copy(bc::Base.Broadcast.Broadcasted{TensorMapStyle}) + lb = TensorAlgebra.tryflattenlinear(bc) + isnothing(lb) && error( "element-wise broadcast is not supported for a `TensorMap`; only linear combinations \ such as `a .+ b` and `2 .* a` are supported" ) + return copy(lb) end # ==================================== pow_diag_safe ====================================== diff --git a/src/linearbroadcasted.jl b/src/linearbroadcasted.jl index 034fcae1..9d4bdeac 100644 --- a/src/linearbroadcasted.jl +++ b/src/linearbroadcasted.jl @@ -25,7 +25,8 @@ abstract type LinearBroadcasted end # Generic interface for LinearBroadcasted subtypes. Base.axes(a::LinearBroadcasted, d::Int) = axes(a)[d] Base.similar(a::LinearBroadcasted) = similar(a, eltype(a)) -Base.similar(a::LinearBroadcasted, elt::Type) = similar(a, elt, axes(a)) +# Forward to 2-arg broadcast `similar` so the result preserves the operands' backend type. +Base.similar(a::LinearBroadcasted, elt::Type) = similar(BC.Broadcasted(a), elt) function Base.show(io::IO, a::LinearBroadcasted) print(io, operation(a), "(", join(arguments(a), ", "), ")") return nothing @@ -42,10 +43,6 @@ function BC.Broadcasted(a::LinearBroadcasted) return BC.Broadcasted(BC.combine_styles(args...), operation(a), args) end -function Base.similar(a::LinearBroadcasted, elt::Type, ax) - return similar(BC.Broadcasted(a), elt, ax) -end - # --- ScaledBroadcasted -------------------------------------------------------- struct ScaledBroadcasted{C <: Number, A} <: LinearBroadcasted