rename, tidy, improve

Author: mcabbott

src/lib/broadcast.jl

@@ -167,25 +167,31 @@ _broadcast(f::F, x...) where F = materialize(broadcasted(f, x...))
 collapse_nothings(xs::AbstractArray{Nothing}) = nothing
 collapse_nothings(xs) = xs
 
-_purefun(::Type{F}) where {F<:Function} = Base.issingletontype(F)
-_purefun(::Type) = false
-_purefun(::Type{typeof(^)}) = false  # fix @testset "power" & @testset "diagonal hessian"
+_dual_purefun(::Type{F}) where {F<:Function} = Base.issingletontype(F)
+_dual_purefun(::Type) = false
+_dual_purefun(::Type{typeof(^)}) = false  # avoid DomainError from negative powers
 
-_dualsafe(x::Numeric{<:Real}) = true
-_dualsafe(x::Ref{<:Numeric{<:Real}}) = true
-_dualsafe(x::Val) = true
-_dualsafe(x::Type) = true
-_dualsafe(x::Symbol) = true
-_dualsafe(x) = false
+_dual_safearg(x::Numeric{<:Real}) = true
+_dual_safearg(x::Ref{<:Numeric{<:Real}}) = true
+_dual_safearg(x::Union{Type,Val,Symbol}) = true  # non-differentiable types
+_dual_safearg(x) = false
+
+# This is Broadcast.combine_eltypes but with dual eltypes:
+_combine_dual_eltypes(f, args::Tuple) =
+  Broadcast.promote_typejoin_union(Base._return_type(f, map(_dual_eltype, args)))
+_dual_eltype(x::Numeric{T}) where {T<:Real} = Dual{Nothing, T, 1} # typeof(Dual(one(T),true))
+_dual_eltype(x) = eltype(x)
 
 @adjoint function broadcasted(::AbstractArrayStyle, f::F, args...) where {F}
-  T = Broadcast.combine_eltypes(f, args)
+  TD = _combine_dual_eltypes(f, args)
   # Avoid generic broadcasting in two easy cases:
-  if T == Bool
+  if TD <: Dual && isconcretetype(TD)
+    if _dual_purefun(F) && all(_dual_safearg, args)
+      y, back = broadcast_forward(f, args...)
+      return y, ȳ -> (nothing, nothing, back(ȳ)...)
+    end
+  elseif TD <: Real && isconcretetype(TD)
     return f.(args...), _->nothing
-  elseif T <: Real && isconcretetype(T) && _purefun(F) && all(_dualsafe, args)
-    y, back = broadcast_forward(f, args...)
-    return y, ȳ -> (nothing, nothing, back(ȳ)...)
   end
   len = inclen(args)
   y∂b = _broadcast((x...) -> _pullback(__context__, f, x...), args...)

test/features.jl

@@ -535,4 +535,10 @@ end
 
   @test gradient(x -> sum(x ./ [1,2,4]), [1,2,pi]) == ([1.0, 0.5, 0.25],)
   @test gradient(x -> sum(map(/, x, [1,2,4])), [1,2,pi]) == ([1.0, 0.5, 0.25],)
+
+  # negative powers
+  @test gradient((x,p) -> sum(x .^ p), [1.0,2.0,4.0], [1,-1,2])[1] ≈ [1.0, -0.25, 8.0]
+  @test gradient((x,p) -> sum(x .^ p), [1.0,2.0,4.0], -1)[1] ≈ [-1.0, -0.25, -0.0625]
+  @test gradient((x,p) -> sum(z -> z^p, x), [1.0,2.0,4.0], -1)[1] ≈ [-1.0, -0.25, -0.0625]
+  @test gradient((x,p) -> mapreduce(z -> z^p, +, x), [1.0,2.0,4.0], -1)[1] ≈ [-1.0, -0.25, -0.0625]
 end