diff --git a/stdlib/LinearAlgebra/src/structuredbroadcast.jl b/stdlib/LinearAlgebra/src/structuredbroadcast.jl index 476d9e3e5cd02..c9c058c867ba0 100644 --- a/stdlib/LinearAlgebra/src/structuredbroadcast.jl +++ b/stdlib/LinearAlgebra/src/structuredbroadcast.jl @@ -78,7 +78,7 @@ find_uplo(bc::Broadcasted) = mapfoldl(find_uplo, merge_uplos, Broadcast.cat_nest function structured_broadcast_alloc(bc, ::Type{Bidiagonal}, ::Type{ElType}, n) where {ElType} uplo = n > 0 ? find_uplo(bc) : 'U' n1 = max(n - 1, 0) - if uplo == 'T' + if count_structedmatrix(Bidiagonal, bc) > 1 && uplo == 'T' return Tridiagonal(Array{ElType}(undef, n1), Array{ElType}(undef, n), Array{ElType}(undef, n1)) end return Bidiagonal(Array{ElType}(undef, n),Array{ElType}(undef, n1), uplo) @@ -135,6 +135,8 @@ iszerodefined(::Type{<:Number}) = true iszerodefined(::Type{<:AbstractArray{T}}) where T = iszerodefined(T) iszerodefined(::Type{<:UniformScaling{T}}) where T = iszerodefined(T) +count_structedmatrix(T, bc::Broadcasted) = sum(Base.Fix2(isa, T), Broadcast.cat_nested(bc); init = 0) + fzeropreserving(bc) = (v = fzero(bc); !ismissing(v) && (iszerodefined(typeof(v)) ? iszero(v) : v == 0)) # 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 diff --git a/stdlib/LinearAlgebra/test/structuredbroadcast.jl b/stdlib/LinearAlgebra/test/structuredbroadcast.jl index 3767fc10055f2..e4e78ad94102c 100644 --- a/stdlib/LinearAlgebra/test/structuredbroadcast.jl +++ b/stdlib/LinearAlgebra/test/structuredbroadcast.jl @@ -96,6 +96,24 @@ using Test, LinearAlgebra @test broadcast!(*, Z, X, Y) == broadcast(*, fX, fY) end end + + @testset "type-stability in Bidiagonal" begin + B2 = @inferred (B -> .- B)(B) + @test B2 isa Bidiagonal + @test B2 == -1 * B + B2 = @inferred (B -> B .* 2)(B) + @test B2 isa Bidiagonal + @test B2 == B + B + B2 = @inferred (B -> 2 .* B)(B) + @test B2 isa Bidiagonal + @test B2 == B + B + B2 = @inferred (B -> B ./ 1)(B) + @test B2 isa Bidiagonal + @test B2 == B + B2 = @inferred (B -> 1 .\ B)(B) + @test B2 isa Bidiagonal + @test B2 == B + end end @testset "broadcast! where the destination is a structured matrix" begin