Make PermutedDims a valid broadcast leaf

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.15.0"
+version = "0.15.1"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/TensorAlgebra.jl

@@ -8,7 +8,7 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc
 if VERSION >= v"1.11.0-DEV.469"
     eval(
         Meta.parse(
-            "public biperm, bipartition, contractopadd!, label_type, matricizeopperm, to_range, zero!, scale!, permuteddims, conjed, ConjArray"
+            "public biperm, bipartition, contractopadd!, label_type, matricizeopperm, to_range, zero!, scale!, permuteddims, PermutedDims, conjed, ConjArray"
         )
     )
 end

src/linearbroadcasted.jl

@@ -281,41 +281,46 @@ linearbroadcasted(::typeof(*), a::Number, b::Number) = a * b
 Per-node trait: can `(f, args...)` be expressed as a `LinearBroadcasted`?
 Extensible by downstream packages for additional linear operations.
 """
+# A leaf operand is treated as an array/tensor by default (`::Any`), and `::Number` marks the
+# scalars. That distinction is what makes the predicate correct: scaling by a scalar is linear
+# while a scalar shift (`a .+ 1`) is affine, and an elementwise array product is nonlinear while
+# scaling is not. Because the array-like slot is `::Any`, any leaf type participates with no
+# change here — it need not be an `AbstractArray` (e.g. a `PermutedDims` wrapper, or a backend
+# tensor); the fold absorbs it via `bipermutedimsopadd!`. An operand that is neither a genuine
+# array-like leaf nor a `Number` (say a scalar buried in an n-ary `+`) is assumed linear here and
+# errors later at the fold rather than silently producing a wrong result.
 islinearbroadcast(f, args...) = false
-islinearbroadcast(::typeof(identity), ::Base.AbstractArrayOrBroadcasted) = true
-islinearbroadcast(::typeof(+), ::Base.AbstractArrayOrBroadcasted...) = true
-islinearbroadcast(::typeof(-), ::Base.AbstractArrayOrBroadcasted) = true
-function islinearbroadcast(
-        ::typeof(-), ::Base.AbstractArrayOrBroadcasted, ::Base.AbstractArrayOrBroadcasted
-    )
-    return true
-end
-islinearbroadcast(::typeof(*), ::Number, ::Base.AbstractArrayOrBroadcasted) = true
-islinearbroadcast(::typeof(\), ::Number, ::Base.AbstractArrayOrBroadcasted) = true
-islinearbroadcast(::typeof(*), ::Base.AbstractArrayOrBroadcasted, ::Number) = true
-islinearbroadcast(::typeof(/), ::Base.AbstractArrayOrBroadcasted, ::Number) = true
-function islinearbroadcast(
-        ::typeof(*), ::Base.AbstractArrayOrBroadcasted, ::Base.AbstractArrayOrBroadcasted
-    )
-    return false
-end
+
+islinearbroadcast(::typeof(identity), ::Any) = true
+islinearbroadcast(::typeof(identity), ::Number) = false
+
+islinearbroadcast(::typeof(+), ::Any...) = true
+islinearbroadcast(::typeof(+), ::Number) = false
+islinearbroadcast(::typeof(+), ::Number, ::Any) = false
+islinearbroadcast(::typeof(+), ::Any, ::Number) = false
+islinearbroadcast(::typeof(+), ::Number, ::Number) = false
+
+islinearbroadcast(::typeof(-), ::Any) = true
+islinearbroadcast(::typeof(-), ::Number) = false
+islinearbroadcast(::typeof(-), ::Any, ::Any) = true
+islinearbroadcast(::typeof(-), ::Number, ::Any) = false
+islinearbroadcast(::typeof(-), ::Any, ::Number) = false
+islinearbroadcast(::typeof(-), ::Number, ::Number) = false
+
+islinearbroadcast(::typeof(*), ::Number, ::Any) = true
+islinearbroadcast(::typeof(*), ::Any, ::Number) = true
+islinearbroadcast(::typeof(*), ::Any, ::Any) = false
 islinearbroadcast(::typeof(*), ::Number, ::Number) = true
-islinearbroadcast(::typeof(conj), ::Base.AbstractArrayOrBroadcasted) = true
-function islinearbroadcast(
-        ::Base.Fix1{typeof(*), <:Number}, ::Base.AbstractArrayOrBroadcasted
-    )
-    return true
-end
-function islinearbroadcast(
-        ::Base.Fix2{typeof(*), <:Number}, ::Base.AbstractArrayOrBroadcasted
-    )
-    return true
-end
-function islinearbroadcast(
-        ::Base.Fix2{typeof(/), <:Number}, ::Base.AbstractArrayOrBroadcasted
-    )
-    return true
-end
+
+islinearbroadcast(::typeof(\), ::Number, ::Any) = true
+islinearbroadcast(::typeof(/), ::Any, ::Number) = true
+
+islinearbroadcast(::typeof(conj), ::Any) = true
+islinearbroadcast(::typeof(conj), ::Number) = false
+
+islinearbroadcast(::Base.Fix1{typeof(*), <:Number}, ::Any) = true
+islinearbroadcast(::Base.Fix2{typeof(*), <:Number}, ::Any) = true
+islinearbroadcast(::Base.Fix2{typeof(/), <:Number}, ::Any) = true
 
 """
     tryflattenlinear(bc::Broadcasted) -> LinearBroadcasted or nothing

src/permutedimsadd.jl

@@ -14,19 +14,32 @@ permuteddims(a, perm) = PermutedDims(a, perm)
 """
     PermutedDims(parent, perm)
 
-Generic lazy permuted-dims wrapper for operands that are not `AbstractArray`s (so
-`Base.PermutedDimsArray` does not apply), modeled on `PermutedDimsArray`: it records `parent`
-and the permutation `perm` without materializing anything. TensorAlgebra never indexes into
-it. It exists so that a downstream extension's `permuteddims` can hand a lazily permuted
-non-array operand to `bipermutedimsopadd!`, whose absorption method composes `perm` into the
-outer bipermutation and forwards to a single leaf call on `parent`.
+Lazy permuted-dims wrapper storing `parent` and the permutation `perm` in fields (unlike
+`Base.PermutedDimsArray`, which encodes `perm` in a type parameter), so it constructs cheaply
+from a runtime permutation. Primarily for internal use to track permutations in linear
+broadcasting.
 """
 struct PermutedDims{P, Perm}
     parent::P
     perm::Perm
 end
 Base.parent(a::PermutedDims) = a.parent
 
+# `PermutedDims` is not an `AbstractArray`, so `Base.Broadcast.broadcastable` would otherwise try
+# to `collect` it. These make it a valid `Broadcasted` leaf (shape = the parent's axes reordered
+# by `perm`); it is never indexed, since the linear-broadcast fold absorbs it into a single
+# `bipermutedimsopadd!` on `parent`. `axes`/`BroadcastStyle` are only reached when the parent is
+# array-like (the only case that broadcasts); the direct `bipermutedimsopadd!` path never uses them.
+Base.ndims(a::PermutedDims) = length(a.perm)
+Base.axes(a::PermutedDims) = map(d -> axes(parent(a), d), a.perm)
+Base.axes(a::PermutedDims, d::Int) = axes(a)[d]
+Base.size(a::PermutedDims) = map(length, axes(a))
+Base.eltype(a::PermutedDims) = eltype(parent(a))
+Base.Broadcast.broadcastable(a::PermutedDims) = a
+function Base.Broadcast.BroadcastStyle(::Type{<:PermutedDims{P}}) where {P}
+    return Base.Broadcast.BroadcastStyle(P)
+end
+
 # ---------------------------------------------------------------------------- #
 # bipermutedimsopadd! — the primary materialization primitive
 # ---------------------------------------------------------------------------- #

test/test_exports.jl

@@ -35,7 +35,8 @@ using Test: @test, @testset
             exports,
             [
                 :biperm, :bipartition, :contractopadd!, :label_type, :matricizeopperm,
-                :to_range, :zero!, :scale!, :permuteddims, :conjed, :ConjArray,
+                :to_range, :zero!, :scale!, :permuteddims, :PermutedDims, :conjed,
+                :ConjArray,
             ]
         )
     end

test/test_linearbroadcasted.jl

@@ -180,4 +180,31 @@ using Test: @test, @test_throws, @testset
         TA.permutedimsopadd!(dest, conj, a, (), 1, 0)
         @test dest[] ≈ conj(a[])
     end
+    @testset "PermutedDims as a linear-broadcast leaf" begin
+        a = randn(3, 4)
+        b = randn(4, 3)
+        pd = TA.PermutedDims(b, (2, 1))   # presents b permuted to shape (3, 4)
+        @test BC.broadcastable(pd) === pd
+        @test axes(pd) == (Base.OneTo(3), Base.OneTo(4))
+        @test eltype(pd) == Float64
+        @test ndims(pd) == 2
+
+        # `a + permuted(b)` flattens (the `PermutedDims` leaf needs no dedicated
+        # `islinearbroadcast` method — it is admitted by the `::Any` leaf slot) and materializes
+        # correctly, the leaf absorbed via `bipermutedimsopadd!` rather than indexed.
+        lb = TA.tryflattenlinear(BC.broadcasted(+, a, pd))
+        @test lb !== nothing
+        dest = similar(a)
+        copyto!(dest, lb)
+        @test dest ≈ a + permutedims(b, (2, 1))
+
+        # `2a - permuted(b)`.
+        lb2 = TA.tryflattenlinear(BC.broadcasted(-, BC.broadcasted(*, 2.0, a), pd))
+        dest2 = similar(a)
+        copyto!(dest2, lb2)
+        @test dest2 ≈ 2a - permutedims(b, (2, 1))
+
+        # A scalar addend is affine, not linear, so it is not flattened.
+        @test TA.tryflattenlinear(BC.broadcasted(+, a, 1)) === nothing
+    end
 end