Preserve structured matrix types more often when broadcasting (#32762)

* Preserve structured matrix types more often when broadcasting

The primary motivation here is exponentiation (fixes #32759). Previously `X .^ 2` would return a Matrix but `two=2; X .^ two` would return the same structured matrix type as `X`. This simply teaches LinearAlgebra a little bit more about broadcast so it can safely compute the zero-preservation property of broadcasts involving `Ref`s, which is sufficient to handle the literal pow argument list.

* fixup negative exponents

Author: mbauman

stdlib/LinearAlgebra/src/structuredbroadcast.jl

@@ -91,14 +91,22 @@ isstructurepreserving(bc::Broadcasted) = isstructurepreserving(bc.f, bc.args...)
 isstructurepreserving(::Union{typeof(abs),typeof(big)}, ::StructuredMatrix) = true
 isstructurepreserving(::TypeFuncs, ::StructuredMatrix) = true
 isstructurepreserving(::TypeFuncs, ::Ref{<:Type}, ::StructuredMatrix) = true
+function isstructurepreserving(::typeof(Base.literal_pow), ::Ref{typeof(^)}, ::StructuredMatrix, ::Ref{Val{N}}) where N
+    return N isa Integer && N > 0
+end
 isstructurepreserving(f, args...) = false
 
 _iszero(n::Number) = iszero(n)
 _iszero(x) = x == 0
 fzeropreserving(bc) = (v = fzero(bc); !ismissing(v) && _iszero(v))
-# Very conservatively only allow Numbers and Types in this speculative zero-test pass
+# Like sparse matrices, we assume that the zero-preservation property of a broadcasted
+# expression is stable.  We can test the zero-preservability by applying the function
+# in cases where all other arguments are known scalars against a zero from the structured
+# matrix. If any non-structured matrix argument is not a known scalar, we give up.
 fzero(x::Number) = x
 fzero(::Type{T}) where T = T
+fzero(r::Ref) = r[]
+fzero(t::Tuple{Any}) = t[1]
 fzero(S::StructuredMatrix) = zero(eltype(S))
 fzero(x) = missing
 function fzero(bc::Broadcast.Broadcasted)

stdlib/LinearAlgebra/test/structuredbroadcast.jl

@@ -28,6 +28,20 @@ using Test, LinearAlgebra
         @test broadcast!(+, Z, fV, fA, X) == broadcast(+, fV, fA, fX)
         @test (Q = broadcast(*, s, fV, fA, X); Q isa Matrix && Q == broadcast(*, s, fV, fA, fX))
         @test broadcast!(*, Z, s, fV, fA, X) == broadcast(*, s, fV, fA, fX)
+
+        @test X .* 2.0 == X .* (2.0,) == fX .* 2.0
+        @test X .* 2.0 isa typeof(X)
+        @test X .* (2.0,) isa typeof(X)
+        @test isequal(X .* Inf, fX .* Inf)
+
+        two = 2
+        @test X .^ 2 ==  X .^ (2,) == fX .^ 2 == X .^ two
+        @test X .^ 2 isa typeof(X)
+        @test X .^ (2,) isa typeof(X)
+        @test X .^ two isa typeof(X)
+        @test X .^ 0 == fX .^ 0
+        @test X .^ -1 == fX .^ -1
+
         for (Y, fY) in zip(structuredarrays, fstructuredarrays)
             @test broadcast(+, X, Y) == broadcast(+, fX, fY)
             @test broadcast!(+, Z, X, Y) == broadcast(+, fX, fY)