diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index cacad7e..3632dfc 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -38,8 +38,14 @@ struct TrackedStyle{N} <: AbstractArrayStyle{N} end (::Type{<:TrackedStyle})(::Val{N}) where {N} = TrackedStyle{N}() Broadcast.BroadcastStyle(::Type{<:TrackedReal}) = TrackedStyle{0}() -# only an element type that fixes `D`, since `getouttype` reads it off -Broadcast.BroadcastStyle(::Type{<:AbstractArray{<:TrackedReal{<:Any, D}, N}}) where {D, N} = +# invariant element types that fix `D`, so that `Union{}` (a subtype of every type) does not match +Broadcast.BroadcastStyle(::Type{<:AbstractArray{TrackedReal{V, D, O}, N}}) where {V, D, O, N} = + TrackedStyle{N}() +Broadcast.BroadcastStyle(::Type{<:AbstractArray{TrackedReal{V, D, O} where {O}, N}}) where {V, D, N} = + TrackedStyle{N}() +Broadcast.BroadcastStyle(::Type{<:AbstractArray{TrackedReal{V, D, O} where {V <: Real}, N}}) where {D, O, N} = + TrackedStyle{N}() +Broadcast.BroadcastStyle(::Type{<:AbstractArray{TrackedReal{V, D, O} where {V <: Real, O}, N}}) where {D, N} = TrackedStyle{N}() # `AbstractArrayStyle{Any}` carries `Any` as its dimension, which `max` cannot compare @@ -101,6 +107,8 @@ end getouttype(::TrackedReal{<:Any, D}) where {D} = D getouttype(::AbstractArray{<:TrackedReal{<:Any, D}}) where {D} = D +# `Union{} <: TrackedReal`, but binds no `D` +getouttype(::AbstractArray{Union{}}) = Union{} getouttype(::Any) = Union{} # a perturbation riding on an argument ends up in the derivative, so `D` has to hold it. An diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 705aec7..22fe8f4 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -56,6 +56,21 @@ end # `Base` styles scalars as `AbstractArrayStyle{0}` @test BroadcastStyle(typeof(track(rand(), tp))) === TrackedStyle{0}() @test TrackedStyle{1}(Val(2)) === TrackedStyle{2}() + # element types that leave the value type or the origin free + @test BroadcastStyle(Vector{TrackedReal{Float64, Float64}}) === TrackedStyle{1}() + @test BroadcastStyle(Vector{TrackedReal{<:Real, Float64}}) === TrackedStyle{1}() + # `TrackedReal{<:Real, Float64, Nothing}`, as built by a comprehension over mixed value types + a = TrackedReal(1.0, 0.0) + b = TrackedReal(ForwardDiff.Dual(1.0, 1.0), 0.0) + @test BroadcastStyle(typeof([x for x in Real[a, b]])) === TrackedStyle{1}() +end + +@testset "arrays that hold no element keep `Base`'s style (#324)" begin + # `Union{} <: TrackedReal` + @test BroadcastStyle(Vector{Union{}}) === DefaultArrayStyle{1}() + @test BroadcastStyle(Matrix{Union{}}) === DefaultArrayStyle{2}() + @test String.(Union{}[]) isa Vector{Union{}} + @test isempty(track(1.0, InstructionTape()) .+ Union{}[]) end @testset "`BroadcastStyle` precedence" begin