fix heterogeneous and abstract cases

Author: stevengj

src/evalpoly.jl

@@ -67,15 +67,19 @@ function evalpoly!(Y::AbstractMatrix, X::AbstractMatrix, p::Union{AbstractVector
     return Y
 end
 
+# fallback cases: call out-of-place _evalpoly
 Base.evalpoly(X::AbstractMatrix, p::Tuple) = _evalpoly(X, p)
 Base.evalpoly(X::AbstractMatrix, ::Tuple{}) = zero(one(X)) # dimensionless zero, i.e. 0 * X^0
 Base.evalpoly(X::AbstractMatrix, p::AbstractVector) = _evalpoly(X, p)
 
-Base.evalpoly(X::StridedMatrix{<:Number}, p::Tuple{Union{Number, UniformScaling}, Vararg{Union{Number, UniformScaling}}}) =
-    evalpoly!(similar(X, Base.promote_op(*, typeof(one(eltype(X))), typeof(_scalarval(p[begin])))), X, p)
-Base.evalpoly(X::StridedMatrix{<:Number}, p::Tuple{AbstractMatrix{<:Number}, Vararg{AbstractMatrix{<:Number}}}) =
-    evalpoly!(similar(X, Base.promote_op(*, typeof(one(eltype(X))), eltype(p[begin]))), X, p)
+# optimized in-place cases, limited to types like homogeneous tuples with length > 1
+# where we can reliably deduce the output type (= type of X * p[2]),
+# and restricted to StridedMatrix (for now) so that we can be more confident that this is a performance win:
+Base.evalpoly(X::StridedMatrix{<:Number}, p::Tuple{T, T, Vararg{T}}) where {T<:Union{Number, UniformScaling}} =
+    evalpoly!(similar(X, Base.promote_op(*, eltype(X), typeof(_scalarval(p[2])))), X, p)
+Base.evalpoly(X::StridedMatrix{<:Number}, p::Tuple{AbstractMatrix{T}, AbstractMatrix{T}, Vararg{AbstractMatrix{T}}}) where {T<:Number} =
+    evalpoly!(similar(X, Base.promote_op(*, eltype(X), T)), X, p)
 Base.evalpoly(X::StridedMatrix{<:Number}, p::AbstractVector{<:Union{Number, UniformScaling}}) =
-    isempty(p) ? _evalpoly(X, p) : evalpoly!(similar(X, Base.promote_op(*, typeof(one(eltype(X))), typeof(_scalarval(p[begin])))), X, p)
+    length(p) < 2 ? _evalpoly(X, p) : evalpoly!(similar(X, Base.promote_op(*, eltype(X), typeof(_scalarval(p[begin+1])))), X, p)
 Base.evalpoly(X::StridedMatrix{<:Number}, p::AbstractVector{<:AbstractMatrix{<:Number}}) =
-    isempty(p) ? _evalpoly(X, p) : evalpoly!(similar(X, Base.promote_op(*, typeof(one(eltype(X))), eltype(p[begin]))), X, p)
+    length(p) < 2 ? _evalpoly(X, p) : evalpoly!(similar(X, Base.promote_op(*, eltype(X), eltype(p[begin+1]))), X, p)

test/generic.jl

@@ -837,23 +837,32 @@ end
     end
 end
 
-using LinearAlgebra: _evalpoly
-naive_evalpoly(X, p) = sum(X^(i-1) * p[i] for i=1:length(p))
+using LinearAlgebra: _evalpoly # fallback routine, which we'll test explicitly
+
+# naive sum, a little complicated since X^0 fails if eltype(X) is abstract:
+naive_evalpoly(X, p) = length(p) == 1 ? one(X) * p[1] : one(X) * p[1] + sum(X^(i-1) * p[i] for i=2:length(p))
 
 @testset "evalpoly" begin
-    for X in ([1 2 3;4 5 6;7 8 9], UpperTriangular([1 2 3;0 5 6;0 0 9]), SymTridiagonal([1,2,3],[4,5]))
+    for X in ([1 2 3;4 5 6;7 8 9], UpperTriangular([1 2 3;0 5 6;0 0 9]),
+              SymTridiagonal([1,2,3],[4,5]), Real[1 2 3;4 5 6;7 8 9])
         @test @inferred(evalpoly(X, ())) == zero(X) == evalpoly(X, Int[])
         @test @inferred(evalpoly(X, (17,))) == one(X) * 17
-        @test @inferred(_evalpoly(X, [1,2,3,4])) == @inferred(evalpoly(X, [1,2,3,4])) ==
-              @inferred(evalpoly(X, (1,2,3,4))) ==
+        @test _evalpoly(X, [1,2,3,4]) == evalpoly(X, [1,2,3,4]) == @inferred(evalpoly(X, (1,2,3,4))) ==
               naive_evalpoly(X, [1,2,3,4]) == 1*one(X) + 2*X + 3X^2 + 4X^3
         @test typeof(evalpoly(X, [1,2,3])) == typeof(evalpoly(X, (1,2,3))) == typeof(_evalpoly(X, [1,2,3])) ==
               typeof(X * X)
 
-        for N in (1,2,4), p in (rand(-10:10, N), UniformScaling.(rand(-10:10, N)), [rand(-5:5,3,3) for _ = 1:N])
+        # _evalpoly is not type-stable if eltype(X) is abstract
+        # because one(Real[...]) returns a Matrix{Int}
+        if isconcretetype(eltype(X))
+            @inferred evalpoly(X, [1,2,3,4])
+            @inferred _evalpoly(X, [1,2,3,4])
+        end
+
+        for N in (1,2,4), p in (Real[1,2], rand(-10:10, N), UniformScaling.(rand(-10:10, N)), [rand(-5:5,3,3) for _ = 1:N])
             @test _evalpoly(X, p) == evalpoly(X, p) == evalpoly(X, Tuple(p)) == naive_evalpoly(X, p)
         end
-        for N in (1,2,4), p in (rand(N), UniformScaling.(rand(N)), [rand(3,3) for _ = 1:N])
+        for N in (1,2,4), p in ((5, 6.7), rand(N), UniformScaling.(rand(N)), [rand(3,3) for _ = 1:N])
             @test _evalpoly(X, p) ≈ evalpoly(X, p) ≈ evalpoly(X, Tuple(p)) ≈ naive_evalpoly(X, p)
         end