tidy, add more lazy cases

Author: mcabbott

src/stage1/broadcast.jl

@@ -29,28 +29,30 @@ function (∂ₙ::∂☆{N})(zc::ZeroBundle{N, typeof(copy)},
   return r
 end
 
-_print(s) = nothing
-# _print(s) = printstyled(s, "\n"; color=:magenta)
+# Reverse mode broadcast rules
+
+using ChainRulesCore: derivatives_given_output
+
+# _print(s) = nothing
+_print(s) = printstyled(s, "\n"; color=:magenta)
 
 # Broadcast over one element is just map
 function (∂⃖ₙ::∂⃖{N})(::typeof(broadcasted), f, a::Array) where {N}
     _print("path 0")
     ∂⃖ₙ(map, f, a)
 end
 
-using ChainRulesCore: derivatives_given_output
-
 (::∂⃖{1})(::typeof(broadcasted), f, args...) = split_bc_rule(f, args...)
 (::∂⃖{1})(::typeof(broadcasted), f, arg::Array) = split_bc_rule(f, arg) # ambiguity
 function split_bc_rule(f::F, args...) where {F}
     T = Broadcast.combine_eltypes(f, args)
-    if T == Bool
-        # Trivial case
+    TΔ = Core.Compiler._return_type(derivatives_given_output, Tuple{T, F, map(eltype, args)...})
+    if eltype(T) == Bool
+        # Trivial case: non-differentiable output
         _print("path 1")
         back_1(_) = ntuple(Returns(ZeroTangent()), length(args)+2)
         return f.(args...), back_1
-    elseif isconcretetype(Core.Compiler._return_type(
-            derivatives_given_output, Tuple{T, F, map(eltype, args)...}))
+    elseif T <: Number && isconcretetype(TΔ)
         # Fast path: just broadcast, and use x & y to find derivative.
         ys = f.(args...)
         _print("path 2")
@@ -65,8 +67,9 @@ function split_bc_rule(f::F, args...) where {F}
         return ys, back_2
     else
         # Slow path: collect all the pullbacks & apply them later.
+        # Since broadcast makes no guarantee about order, this does not bother to try to reverse it.
         _print("path 3")
-        ys, backs = splitcast(rrule_via_ad, DiffractorRuleConfig(), f, args...)
+        ys, backs = splitcast(∂⃖{1}(), f, args...)
         function back_3(dys)
             deltas = splitmap(backs, unthunk(dys)) do back, dy
                 map(unthunk, back(dy))
@@ -78,8 +81,11 @@ function split_bc_rule(f::F, args...) where {F}
     end
 end
 
+# This uses "mulltimap"-like constructs:
+
 using StructArrays
-splitmap(f, args...) = StructArrays.components(StructArray(Iterators.map(f, args...))) # warning: splitmap(identity, [1,2,3,4]) === NamedTuple()
+splitmap(f, args...) = StructArrays.components(StructArray(Iterators.map(f, args...)))
+# warning: splitmap(identity, [1,2,3,4]) === NamedTuple()
 splitcast(f, args...) = StructArrays.components(StructArray(Broadcast.instantiate(Broadcast.broadcasted(f, args...))))
 
 # For certain cheap operations we can easily allow fused broadcast:
@@ -107,7 +113,7 @@ end
 
 using LinearAlgebra: dot
 
-function (::∂⃖{1})(::typeof(broadcasted), ::typeof(*), x, y)
+function (::∂⃖{1})(::typeof(broadcasted), ::typeof(*), x, y)  # should this be vararg, or will laziness handle it?
     broadcasted(*, x, y), Δ -> let Δun = unthunk(Δ)
         _print("broadcast *")
         dx = eltype(x)==Bool ? NoTangent() : x isa Number ? dot(y, Δun) : unbroadcast(x, Δun .* conj.(y))
@@ -117,41 +123,88 @@ function (::∂⃖{1})(::typeof(broadcasted), ::typeof(*), x, y)
         (NoTangent(), NoTangent(), dx, dy)
     end
 end
+# Alternative to `x isa Number` etc above... but not quite right!
+# (::∂⃖{1})(::typeof(broadcasted), ::typeof(*), x, y::Number) = rrule_via_ad(DiffractorRuleConfig(), *, x, y)
+
+function (::∂⃖{1})(::typeof(broadcasted), ::typeof(Base.literal_pow), ::typeof(^), x, ::Val{2})
+    _print("broadcast ^2")
+    broadcasted(*, x, x), Δ -> begin
+        dx = unbroadcast(x, 2 .* Δ .* conj.(x))
+        (NoTangent(), NoTangent(), NoTangent(), dx, NoTangent()) 
+    end
+end
+function (::∂⃖{1})(::typeof(broadcasted), ::typeof(Base.literal_pow), ::typeof(^), x::Number, ::Val{2})
+    _print("simple ^2")
+    x^2, Δ -> (NoTangent(), NoTangent(), NoTangent(), 2 * Δ * conj(x), NoTangent())
+end
+
+# function (::∂⃖{1})(::typeof(broadcasted), ::typeof(/), x, y) # not obvious whether this is better than automatic
+#     broadcasted(/, x, y), Δ -> let Δun = unthunk(Δ)
+#         _print("broadcast /")
+#         dx = unbroadcast(x, Δ ./ conj.(y))
+#         dy = unbroadcast(y, .-Δ .* conj.(res ./ y))
+#         (NoTangent(), NoTangent(), dx, dy)
+#     end
+# end
+function (::∂⃖{1})(::typeof(broadcasted), ::typeof(/), x, y::Number)
+    _print("simple /")
+    z, back = ∂⃖{1}()(/, x, y)
+    z, Δ -> begin
+        _, dx, dy = back(Δ)
+        (NoTangent(), NoTangent(), dx, dy)  # maybe there should be a funciton for this? Use for conj, identity too
+    end
+end
 
 (::∂⃖{1})(::typeof(broadcasted), ::typeof(conj), x) =
     broadcasted(conj, x), Δ -> (NoTangent(), conj(unthunk(Δ)))
 (::∂⃖{1})(::typeof(broadcasted), ::typeof(conj), x::AbstractArray{Real}) =
     x, Δ -> (NoTangent(), Δ)
 
+(::∂⃖{1})(::typeof(broadcasted), ::typeof(identity), x) =
+    x, Δ -> (NoTangent(), Δ)
+
+# All broadcasts use `unbroadcast` to reduce to correct shape:
+
 function unbroadcast(x::Base.AbstractArrayOrBroadcasted, dx)
     N = ndims(dx)
     if length(x) == length(dx)
         ProjectTo(x)(dx)  # handles trivial reshapes, offsets, structured matrices, row vectors
     else
-        # This is an awful hack to get type-stable `dims`
-        dims = ntuple(d -> get(size(x), d, 1) == 1 ? d : N+1, N)
+        dims = ntuple(d -> get(size(x), d, 1) == 1 ? d : N+1, N)  # awful hack to get type-stable `dims`
         ProjectTo(x)(sum(dx; dims))
     end
 end
 unbroadcast(x::Base.AbstractArrayOrBroadcasted, dx::NoTangent) = NoTangent()
 
+unbroadcast(x::T, dx) where {T<:Tuple{Any}} = ProjectTo(x)(Tangent{T}(sum(dx)))
+function unbroadcast(x::T, dx) where {T<:Tuple{Vararg{Any,N}}} where {N}
+    _print("unbroadcast tuple")
+    val = if length(x) == length(dx)
+        dx
+    else
+        sum(dx; dims=2:ndims(dx))
+    end
+    ProjectTo(x)(NTuple{length(x)}(val)) # Tangent
+end
+
+unbroadcast(f::Function, df) = sum(df)
 unbroadcast(x::Number, dx) = ProjectTo(x)(sum(dx))
-unbroadcast(f::Function, df) = ProjectTo(x)(sum(df))
 unbroadcast(x::Base.RefValue, dx) = ProjectTo(x)(Ref(sum(dx)))
 
 unbroadcast(::Bool, dx) = NoTangent()
 unbroadcast(::AbstractArray{Bool}, dx) = NoTangent()
 unbroadcast(::AbstractArray{Bool}, ::NoTangent) = NoTangent()  # ambiguity
 unbroadcast(::Val, dx) = NoTangent()
-# Maybe more non-diff types? Some fallback?
 
-unbroadcast(x::T, dx) where {T<:Tuple{Any}} = ProjectTo(x)(Tangent{T}(sum(dx)))
-function unbroadcast(x::T, dx) where {T<:Tuple{Vararg{Any,N}}} where {N}
-    _print("unbroadcast tuple")
-    val = if length(x) == length(dx)
-        dx
+function unbroadcast(x, dx)
+    p = ProjectTo(x)
+    if dx isa AbstractZero || p isa ProjectTo{<:AbstractZero}
+        return NoTangent()
+    end
+    b = Broadcast.broadcastable(x)
+    if b isa Ref  # then x is scalar under broadcast
+        return p(sum(dx))
     else
-        sum(dx; dims=2:ndims(dx))
+        error("don't know how to handle broadcast gradient for x::$(typeof(x))")
     end
-    ProjectTo(x)(NTuple{length(x)}(val)) # Tangent
 end

test/runtests.jl

@@ -216,6 +216,7 @@ z45, delta45 = frule_via_ad(DiffractorRuleConfig(), (0,1), x -> log(exp(x)), 2)
 
 # Broadcasting
 @test gradient(x -> sum(x ./ x), [1,2,3]) == ([0,0,0],)  # derivatives_given_output
+@test gradient(x -> sum(sqrt.(atan.(x, x'))), [1,2,3])[1] ≈ [0.2338, -0.0177, -0.0661] atol=1e-3
 @test gradient(x -> sum(exp.(log.(x))), [1,2,3]) == ([1,1,1],)
 
 @test_broken gradient(x -> sum((exp∘log).(x)), [1,2,3]) == ([1,1,1],)  # stores pullback
@@ -229,6 +230,9 @@ exp_log(x) = exp(log(x))
 @test gradient(x -> sum(sum, Ref(x) ./ x), [1,2,3])[1] ≈ [-4.1666, 0.3333, 1.1666] atol=1e-3
 @test gradient(x -> sum(sum, (x,) ./ x), [1,2,3])[1] ≈ [-4.1666, 0.3333, 1.1666] atol=1e-3
 
+@test unthunk.(gradient(x -> sum(x ./ 4), [1,2,3])) == ([0.25, 0.25, 0.25],)
+@test gradient(x -> sum([1,2,3] ./ x), 4) == (-0.375,)
+
 @test gradient(x -> sum(x .> 2), [1,2,3]) == (ZeroTangent(),)  # Bool output
 @test gradient(x -> sum(1 .+ iseven.(x)), [1,2,3]) == (ZeroTangent(),)
 @test gradient((x,y) -> sum(x .== y), [1,2,3], [1 2 3]) == (ZeroTangent(), ZeroTangent())