Support linear broadcasting over TensorMaps

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 <support@itensor.org> and contributors"]
 
 [workspace]

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  ======================================

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