From aacb2f2572b208f64c81dde8475026526b9602b9 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 1 Jul 2026 18:58:01 -0400 Subject: [PATCH 1/2] Make PermutedDims a valid broadcast leaf ## Summary Makes `PermutedDims` a valid `Base.Broadcasted` leaf, so a linear-combination broadcast that carries a lazily permuted operand (adding a permuted array, say) flattens through `tryflattenlinear` and materializes via `bipermutedimsopadd!` on the parent, instead of hitting the `broadcastable` fallback that would `collect` it. It implements the broadcast-leaf interface (`axes`, `eltype`, `broadcastable`, `BroadcastStyle`) and is still never indexed, since the fold absorbs it into a single leaf call. Generalizes `islinearbroadcast` so any leaf type participates, not just `AbstractArray`/`Broadcasted`. The default leaf slot is now `::Any`, with `::Number` marking the scalars, and that distinction is what keeps 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. A `PermutedDims` wrapper (or a backend tensor) is then admitted by the `::Any` slot with no dedicated method, and `PermutedDims` becomes public. Builds on https://github.com/ITensor/TensorAlgebra.jl/pull/195. --- Project.toml | 2 +- src/TensorAlgebra.jl | 2 +- src/linearbroadcasted.jl | 71 ++++++++++++++++++---------------- src/permutedimsadd.jl | 29 +++++++++++--- test/test_exports.jl | 3 +- test/test_linearbroadcasted.jl | 27 +++++++++++++ 6 files changed, 92 insertions(+), 42 deletions(-) diff --git a/Project.toml b/Project.toml index 31e1b69..91ca84e 100644 --- a/Project.toml +++ b/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 and contributors"] [workspace] diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index f9d5366..92bbcd2 100644 --- a/src/TensorAlgebra.jl +++ b/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 diff --git a/src/linearbroadcasted.jl b/src/linearbroadcasted.jl index 88440ec..125419f 100644 --- a/src/linearbroadcasted.jl +++ b/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 diff --git a/src/permutedimsadd.jl b/src/permutedimsadd.jl index 82d1780..843cf81 100644 --- a/src/permutedimsadd.jl +++ b/src/permutedimsadd.jl @@ -14,12 +14,14 @@ 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. It is deliberately not an `AbstractArray`, so it can wrap operands +that are not arrays (e.g. a backend tensor) as well as arrays. It is a valid `Base.Broadcasted` +leaf: a linear-combination broadcast that carries it (adding a permuted operand, say) flattens +through `tryflattenlinear` and is materialized by `bipermutedimsopadd!`, whose absorption method +composes `perm` into the outer bipermutation and forwards a single leaf call on `parent`. +TensorAlgebra never indexes into it. """ struct PermutedDims{P, Perm} parent::P @@ -27,6 +29,21 @@ struct PermutedDims{P, 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 # ---------------------------------------------------------------------------- # diff --git a/test/test_exports.jl b/test/test_exports.jl index 1ccdd1c..2f39476 100644 --- a/test/test_exports.jl +++ b/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 diff --git a/test/test_linearbroadcasted.jl b/test/test_linearbroadcasted.jl index 878b877..5535f9a 100644 --- a/test/test_linearbroadcasted.jl +++ b/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 From e04423de60428e95a06aac2ea9a7b742f42b1d87 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 1 Jul 2026 19:07:29 -0400 Subject: [PATCH 2/2] Trim the PermutedDims docstring --- src/permutedimsadd.jl | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/permutedimsadd.jl b/src/permutedimsadd.jl index 843cf81..e0a6478 100644 --- a/src/permutedimsadd.jl +++ b/src/permutedimsadd.jl @@ -16,12 +16,8 @@ permuteddims(a, perm) = PermutedDims(a, perm) 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. It is deliberately not an `AbstractArray`, so it can wrap operands -that are not arrays (e.g. a backend tensor) as well as arrays. It is a valid `Base.Broadcasted` -leaf: a linear-combination broadcast that carries it (adding a permuted operand, say) flattens -through `tryflattenlinear` and is materialized by `bipermutedimsopadd!`, whose absorption method -composes `perm` into the outer bipermutation and forwards a single leaf call on `parent`. -TensorAlgebra never indexes into it. +from a runtime permutation. Primarily for internal use to track permutations in linear +broadcasting. """ struct PermutedDims{P, Perm} parent::P