fixup

Author: mcabbott

src/extra_rules.jl

@@ -138,27 +138,17 @@ struct NonDiffOdd{N, O, P}; end
 (::NonDiffEven{N, O, O})(Δ...) where {N, O} = error()
 
 # WARNING: Method definition rrule(typeof(Core.apply_type), Any, Any...) in module ChainRules at /Users/me/.julia/packages/ChainRules/kkDLd/src/rulesets/Core/core.jl:10 overwritten in module Diffractor at /Users/me/.julia/dev/Diffractor/src/extra_rules.jl:140.
-# @Base.pure function ChainRulesCore.rrule(::typeof(Core.apply_type), head, args...)
-#     Core.apply_type(head, args...), NonDiffOdd{plus1(plus1(length(args))), 1, 1}()
-# end
+@Base.pure function ChainRulesCore.rrule(::typeof(Core.apply_type), head, args...)
+    Core.apply_type(head, args...), NonDiffOdd{plus1(plus1(length(args))), 1, 1}()
+end
 
 function ChainRulesCore.rrule(::typeof(Core.tuple), args...)
     Core.tuple(args...), Δ->Core.tuple(NoTangent(), Δ...)
 end
 
-# TODO: What to do about these integer rules
-# @ChainRulesCore.non_differentiable Base.rem(a::Integer, b::Type)  # now in CR 1.18
-
 ChainRulesCore.canonicalize(::ChainRulesCore.ZeroTangent) = ChainRulesCore.ZeroTangent()
 ChainRulesCore.canonicalize(::NoTangent) = NoTangent()
 
-# # Skip AD'ing through the axis computation
-# function ChainRules.rrule(::typeof(Base.Broadcast.instantiate), bc::Base.Broadcast.Broadcasted)
-#     return Base.Broadcast.instantiate(bc), Δ->begin
-#         Core.tuple(NoTangent(), Δ)
-#     end
-# end
-
 
 using StaticArrays
 
@@ -201,22 +191,6 @@ function ChainRules.rrule(::typeof(map), ::typeof(+), A::AbstractVector, B::Abst
     map(+, A, B), Δ->(NoTangent(), NoTangent(), Δ, Δ)
 end
 
-# https://github.com/JuliaDiff/ChainRules.jl/blob/main/src/rulesets/Base/array.jl#L7
-# function ChainRules.rrule(AT::Type{<:Array{T,N}}, x::AbstractArray{S,N}) where {T,S,N}
-#     # We're leaving these in the eltype that the cotangent vector already has.
-#     # There isn't really a good reason to believe we should convert to the
-#     # original array type, so don't unless explicitly requested.
-#     AT(x), Δ->(NoTangent(), Δ)
-# end
-
-# WARNING: Method definition rrule(Type{var"#s260"} where var"#s260"<:(Array{T, N} where N where T), UndefInitializer, Any...) in module ChainRules at /Users/me/.julia/packages/ChainRules/kkDLd/src/rulesets/Base/array.jl:5 overwritten in module Diffractor at /Users/me/.julia/dev/Diffractor/src/extra_rules.jl:209.
-# function ChainRules.rrule(AT::Type{<:Array}, undef::UndefInitializer, args...)
-#     # We're leaving these in the eltype that the cotangent vector already has.
-#     # There isn't really a good reason to believe we should convert to the
-#     # original array type, so don't unless explicitly requested.
-#     AT(undef, args...), Δ->(NoTangent(), NoTangent(), ntuple(_->NoTangent(), length(args))...)
-# end
-
 function unzip_tuple(t::Tuple)
     map(x->x[1], t), map(x->x[2], t)
 end
@@ -256,7 +230,6 @@ function ChainRules.frule(_, ::Type{Vector{T}}, undef::UndefInitializer, dims::I
     Vector{T}(undef, dims...), zeros(T, dims...)
 end
 
-# @ChainRules.non_differentiable Base.:(|)(a::Integer, b::Integer) CR#558
 @ChainRules.non_differentiable Base.throw(err)
 @ChainRules.non_differentiable Core.Compiler.return_type(args...)
 
@@ -273,3 +246,13 @@ end
 # ERROR: ArgumentError: Tangent for the primal Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}} should be backed by a AbstractDict type, not by NamedTuple{(:data,), Tuple{ChainRulesCore.ZeroTangent}}.
 ChainRulesCore._backing_error(::Type{<:Base.Pairs{Symbol}}, ::Type{<:NamedTuple}, _) = nothing  # solves that!
 
+# Rather than have a rule for broadcasted 3-arg *, just send it to the efficient path:
+ChainRulesCore.derivatives_given_output(Ω, ::typeof(*), x::Number, y::Number, z::Number) = ((y*z, x*z, x*y),)
+function ChainRulesCore.derivatives_given_output(Ω, ::typeof(*), x::Number, y::Number, z::Number, w::Number)
+    xy = x*y
+    zw = z*w
+    ((y*zw, x*zw, xy*w, xy*z),)
+end
+
+# Fixes @test (x->sum(isa_control_flow(Matrix{Float64}, x)))'(Float32[1 2;]) == [1.0 1.0;]
+(project::ProjectTo{<:AbstractArray})(th::InplaceableThunk) = project(unthunk(th))

src/stage1/broadcast.jl

@@ -34,12 +34,14 @@ end
 using ChainRulesCore: derivatives_given_output
 
 # Broadcast over one element is just map
-function (∂⃖ₙ::∂⃖{N})(::typeof(broadcasted), f, a::Array) where {N}
-    ∂⃖ₙ(map, f, a)
-end
+# function (∂⃖ₙ::∂⃖{N})(::typeof(broadcasted), f, a::Array) where {N}
+#     ∂⃖ₙ(map, f, a)
+# end
+
+(::∂⃖{1})(::typeof(copy), bc::Broadcast.Broadcasted) = copy(bc), Δ -> (NoTangent(), Δ)
 
-(::∂⃖{1})(::typeof(broadcasted), f, args...) = split_bc_rule(f, args...)
-(::∂⃖{1})(::typeof(broadcasted), f, arg::Array) = split_bc_rule(f, arg) # ambiguity
+(::∂⃖{1})(::typeof(broadcasted), f::F, args...) where {F} = split_bc_rule(f, args...)
+# (::∂⃖{1})(::typeof(broadcasted), f::F, arg::Array) where {F} = split_bc_rule(f, arg) # ambiguity
 function split_bc_rule(f::F, args::Vararg{Any,N}) where {F,N}
     T = Broadcast.combine_eltypes(f, args)
     TΔ = Core.Compiler._return_type(derivatives_given_output, Tuple{T, F, map(eltype, args)...})
@@ -48,17 +50,17 @@ function split_bc_rule(f::F, args::Vararg{Any,N}) where {F,N}
         back_1(_) = ntuple(Returns(ZeroTangent()), length(args)+2)
         return f.(args...), back_1
     elseif T <: Number && isconcretetype(TΔ)
-        # Fast path: just broadcast, and use x & y to find derivative.
+        # Fast path: just broadcast, and use arguments & result to find derivatives.
         ys = f.(args...)
         function back_2_one(dys)  # For f.(x) we do not need StructArrays / unzip at all
             delta = broadcast(unthunk(dys), ys, args...) do dy, y, a
                 das = only(derivatives_given_output(y, f, a))
-                dy * conj(only(das))
+                dy * conj(only(das))  # possibly this * should be made nan-safe.
             end
             (NoTangent(), NoTangent(), unbroadcast(only(args), delta))
         end
         function back_2_many(dys)
-            deltas = splitcast(unthunk(dys), ys, args...) do dy, y, as...
+            deltas = tuplecast(unthunk(dys), ys, args...) do dy, y, as...
                 das = only(derivatives_given_output(y, f, as...))
                 map(da -> dy * conj(da), das)
             end
@@ -70,62 +72,76 @@ function split_bc_rule(f::F, args::Vararg{Any,N}) where {F,N}
         # Slow path: collect all the pullbacks & apply them later.
         # (Since broadcast makes no guarantee about order of calls, and un-fusing 
         # can change the number of calls, this does not bother to try to reverse.)
-        ys, backs = splitcast(∂⃖{1}(), f, args...)
+        ys3, backs = tuplecast(∂⃖{1}(), f, args...)
         function back_3(dys)
-            deltas = splitmap(backs, unthunk(dys)) do back, dy
+            deltas = tuplecast(backs, unthunk(dys)) do back, dy  # could be map, sizes match
                 map(unthunk, back(dy))
             end
-            dargs = map(unbroadcast, args, Base.tail(deltas))  # no real need to close over args here
+            dargs = map(unbroadcast, args, Base.tail(deltas))
             (NoTangent(), sum(first(deltas)), dargs...)
         end
         back_3(::AbstractZero) = (NoTangent(), map(Returns(ZeroTangent()), args)...)
-        return ys, back_3
+        return ys3, back_3
     end
 end
 
+# Don't run broadcasting on scalars
+function split_bc_rule(f::F, args::Number...) where {F}
+    z, back = ∂⃖{1}()(f, args...)
+    z, dz -> (NoTangent(), back(dz)...)
+end
+
+split_bc_rule(::typeof(identity), x) = x, Δ -> (NoTangent(), NoTangent(), Δ)
+split_bc_rule(::typeof(identity), x::Number) = x, Δ -> (NoTangent(), NoTangent(), Δ)
+
 # Skip AD'ing through the axis computation
 function (::∂⃖{1})(::typeof(Base.Broadcast.instantiate), bc::Base.Broadcast.Broadcasted)
     uninstantiate(Δ) = Core.tuple(NoTangent(), Δ)
     return Base.Broadcast.instantiate(bc), uninstantiate
 end
 
-# This uses "multimap"-like constructs:
 using StructArrays
-splitmap(f, args...) = StructArrays.components(StructArray(Iterators.map(f, args...)))
-splitcast(f, args...) = StructArrays.components(StructArray(Broadcast.instantiate(Broadcast.broadcasted(f, args...))))
+
+function tuplecast(f::F, args...) where {F}
+    T = Broadcast.combine_eltypes(f, args)
+    if isconcretetype(T)
+        T <: Tuple || throw(ArgumentError("tuplecast(f, args) only works on functions returning a tuple."))
+    end
+    bc = Broadcast.instantiate(Broadcast.broadcasted(f, args...))
+    StructArrays.components(StructArray(bc))
+end
 
 # For certain cheap operations we can easily allow fused broadcast:
+const NumericOrBroadcast = Union{Number, AbstractArray{<:Number}, NTuple{<:Any,Number}, Broadcast.Broadcasted}
 
-(::∂⃖{1})(::typeof(broadcasted), ::typeof(+), args...) = lazy_bc_plus(args...)
-(::∂⃖{1})(::typeof(broadcasted), ::typeof(+), arg::Array) = lazy_bc_plus(arg) # ambiguity
+(::∂⃖{1})(::typeof(broadcasted), ::typeof(+), args::NumericOrBroadcast...) = lazy_bc_plus(args...)
+(::∂⃖{1})(::typeof(broadcasted), ::typeof(+), args::Number) = split_bc_rule(+, args...)
 function lazy_bc_plus(xs...) where {F}
     broadcasted(+, xs...), Δraw -> let Δ = unthunk(Δraw)
         (NoTangent(), NoTangent(), map(x -> unbroadcast(x, Δ), xs)...)
     end
 end
 
-(::∂⃖{1})(::typeof(copy), bc::Broadcast.Broadcasted) = copy(bc), Δ -> (NoTangent(), Δ)
-
-function (::∂⃖{1})(::typeof(broadcasted), ::typeof(-), x, y)
+(::∂⃖{1})(::typeof(broadcasted), ::typeof(-), x::Number, y::Number) = split_bc_rule(-, x, y)
+function (::∂⃖{1})(::typeof(broadcasted), ::typeof(-), x::NumericOrBroadcast, y::NumericOrBroadcast)
     broadcasted(-, x, y), Δraw -> let Δ = unthunk(Δraw)
         (NoTangent(), NoTangent(), unbroadcast(x, Δ), -unbroadcast(y, Δ))
-        # Ideally you could fuse the - into unbroadcast, mapreduce() not sum, when y is a smaller array
     end
 end
 
 using LinearAlgebra: dot
-const Numeric{T<:Number} = Union{T, AbstractArray{T}}
 
-function (::∂⃖{1})(::typeof(broadcasted), ::typeof(*), x::Numeric, y::Numeric)
+(::∂⃖{1})(::typeof(broadcasted), ::typeof(*), x::Number, y::Number) = split_bc_rule(*, x, y)
+function (::∂⃖{1})(::typeof(broadcasted), ::typeof(*), x::NumericOrBroadcast, y::NumericOrBroadcast)
     broadcasted(*, x, y), Δraw -> let Δ = unthunk(Δraw)
-        dx = eltype(x)==Bool ? NoTangent() : x isa Number ? dot(y, Δ) : unbroadcast(x, Δ .* conj.(y))
-        dy = eltype(y)==Bool ? NoTangent() : y isa Number ? dot(x, Δ) : unbroadcast(y, Δ .* conj.(x))
-        # When x is an array but a smaller one, instead of dot you may be able to use mapreduce()
-        (NoTangent(), NoTangent(), dx, dy)
+        (NoTangent(), NoTangent(), _back_star(x, y, Δ), _back_star(y, x, Δ))
     end
 end
+_back_star(x, y, Δ) = unbroadcast(x, Δ .* conj.(y))
+_back_star(x::Number, y, Δ) = dot(y, Δ)
+_back_star(x::Bool, y, Δ) = NoTangent()
 
-function (::∂⃖{1})(::typeof(broadcasted), ::typeof(Base.literal_pow), ::typeof(^), x, ::Val{2})
+function (::∂⃖{1})(::typeof(broadcasted), ::typeof(Base.literal_pow), ::typeof(^), x::NumericOrBroadcast, ::Val{2})
     broadcasted(*, x, x), Δ -> begin
         dx = unbroadcast(x, 2 .* unthunk(Δ) .* conj.(x))
         (NoTangent(), NoTangent(), NoTangent(), dx, NoTangent()) 
@@ -135,41 +151,40 @@ function (::∂⃖{1})(::typeof(broadcasted), ::typeof(Base.literal_pow), ::type
     x^2, Δ -> (NoTangent(), NoTangent(), NoTangent(), 2 * Δ * conj(x), NoTangent())
 end
 
-function (::∂⃖{1})(::typeof(broadcasted), ::typeof(/), x::Numeric, y::Number)
-    z, back = ∂⃖{1}()(/, x, y)
-    z, dz -> begin
-        _, dx, dy = back(dz)
+(::∂⃖{1})(::typeof(broadcasted), ::typeof(/), x::Number, y::Number) = split_bc_rule(/, x, y)
+function (::∂⃖{1})(::typeof(broadcasted), ::typeof(/), x::NumericOrBroadcast, y::Number)
+    z = broadcast(/, x, y)
+    z, Δth -> let Δ = unthunk(Δth)
+        dx = unbroadcast(x, Δ ./ conj.(y))
+        dy = -dot(z, Δ) / (conj(y))  # the reason to be eager is to allow dot here
         (NoTangent(), NoTangent(), dx, dy)
     end
 end
 
-(::∂⃖{1})(::typeof(broadcasted), ::typeof(identity), x) = x, identity_pullback
-(::∂⃖{1})(::typeof(broadcasted), ::typeof(identity), x::Array) = x, identity_pullback # ambiguity
-identity_pullback(Δ) = (NoTangent(), NoTangent(), Δ)
+(::∂⃖{1})(::typeof(broadcasted), ::typeof(identity), x) = split_bc_rule(identity, x)
+# (::∂⃖{1})(::typeof(broadcasted), ::typeof(identity), x::Array) = split_bc_rule(identity, x) # ambiguity
 
-(::∂⃖{1})(::typeof(broadcasted), ::typeof(conj), x::AbstractArray{Real}) = x, identity_pullback
-(::∂⃖{1})(::typeof(broadcasted), ::typeof(conj), x::Array{Real}) = x, identity_pullback
+(::∂⃖{1})(::typeof(broadcasted), ::typeof(conj), x::AbstractArray{Real}) = split_bc_rule(identity, x)
+# (::∂⃖{1})(::typeof(broadcasted), ::typeof(conj), x::Array{Real}) = split_bc_rule(identity, x)  # ambiguity
 (::∂⃖{1})(::typeof(broadcasted), ::typeof(conj), x) =
     broadcasted(conj, x), Δ -> (NoTangent(), conj(unthunk(Δ)))
 (::∂⃖{1})(::typeof(broadcasted), ::typeof(conj), x::Array) =
     broadcasted(conj, x), Δ -> (NoTangent(), conj(unthunk(Δ)))
 
-# All broadcasts use `unbroadcast` to reduce to correct shape:
-
+# Reverse mode broadcasting uses `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
-        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))
+        dims = ntuple(d -> get(size(x), d, 1) == 1 ? d : N+1, N)  # hack to get type-stable `dims`
+        ProjectTo(x)(sum(dx; dims))  # ideally this sum might be thunked?
     end
 end
 unbroadcast(x::Base.AbstractArrayOrBroadcasted, dx::AbstractZero) = dx
 
 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