update, rm comments

Author: mcabbott

Project.toml

@@ -8,13 +8,14 @@ ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 
 [compat]
-ChainRules = "1.5"
-ChainRulesCore = "1.4"
+ChainRules = "1.17"
+ChainRulesCore = "1.11"
 Combinatorics = "1"
 StaticArrays = "1"
 StatsBase = "0.33"

src/extra_rules.jl

@@ -137,9 +137,10 @@ struct NonDiffOdd{N, O, P}; end
 # This should not happen
 (::NonDiffEven{N, O, O})(Δ...) where {N, O} = error()
 
-@Base.pure function ChainRulesCore.rrule(::typeof(Core.apply_type), head, args...)
-    Core.apply_type(head, args...), NonDiffOdd{plus1(plus1(length(args))), 1, 1}()
-end
+# 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
 
 function ChainRulesCore.rrule(::typeof(Core.tuple), args...)
     Core.tuple(args...), Δ->Core.tuple(NoTangent(), Δ...)
@@ -206,12 +207,13 @@ function ChainRules.rrule(AT::Type{<:Array{T,N}}, x::AbstractArray{S,N}) where {
     AT(x), Δ->(NoTangent(), Δ)
 end
 
-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
+# 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)

src/stage1/broadcast.jl

@@ -29,8 +29,12 @@ function (∂ₙ::∂☆{N})(zc::ZeroBundle{N, typeof(copy)},
   return r
 end
 
+_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
 
@@ -40,16 +44,16 @@ using ChainRulesCore: derivatives_given_output
 (::∂⃖{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 && Base.issingletontype(F)
+    if T == Bool
         # Trivial case
+        _print("path 1")
         back_1(_) = ntuple(Returns(ZeroTangent()), length(args)+2)
         return f.(args...), back_1
-    # elseif all(a -> a isa Numeric, args) && isconcretetype(Core.Compiler._return_type(
     elseif isconcretetype(Core.Compiler._return_type(
             derivatives_given_output, Tuple{T, F, map(eltype, args)...}))
         # Fast path: just broadcast, and use x & y to find derivative.
         ys = f.(args...)
-        # println("2")
+        _print("path 2")
         function back_2(dys)
             deltas = splitcast(unthunk(dys), ys, args...) do dy, y, as...
                 das = only(derivatives_given_output(y, f, as...))
@@ -61,7 +65,7 @@ function split_bc_rule(f::F, args...) where {F}
         return ys, back_2
     else
         # Slow path: collect all the pullbacks & apply them later.
-        # println("3")
+        _print("path 3")
         ys, backs = splitcast(rrule_via_ad, DiffractorRuleConfig(), f, args...)
         function back_3(dys)
             deltas = splitmap(backs, unthunk(dys)) do back, dy
@@ -78,108 +82,13 @@ using StructArrays
 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...))))
 
-unbroadcast(f::Function, x̄) = accum_sum(x̄)
-unbroadcast(::Val, _) = NoTangent()
-unbroadcast(x::AbstractArray, x̄::NoTangent) = NoTangent()
-accum_sum(xs::AbstractArray{<:NoTangent}; dims = :) = NoTangent()
-
-#=
-
-julia> xs = randn(10_000);
-julia> @btime Zygote.gradient(x -> sum(abs2, x), $xs)
-  4.744 μs (2 allocations: 78.17 KiB)
-julia> @btime Diffractor.unthunk.(gradient(x -> sum(abs2, x), $xs));
-  3.307 μs (2 allocations: 78.17 KiB)
-
-# Simple function
-
-julia> @btime Zygote.gradient(x -> sum(abs2, exp.(x)), $xs);
-  72.541 μs (29 allocations: 391.47 KiB)     # with dual numbers -- like 4 copies
-
-julia> @btime gradient(x -> sum(abs2, exp.(x)), $xs);
-  45.875 μs (36 allocations: 235.47 KiB)  # fast path -- one copy forward, one back
-  44.042 μs (32 allocations: 313.48 KiB)  # slow path -- 3 copies, extra is closure? 
-  61.167 μs (12 allocations: 703.41 KiB)  # with `map` rule as before -- worse
-
-
-# Composed function, Zygote struggles
-
-julia> @btime Zygote.gradient(x -> sum(abs2, (identity∘cbrt).(x)), $xs);
-  97.167 μs (29 allocations: 391.61 KiB)     # with dual numbers (Zygote master)
-  93.238 ms (849567 allocations: 19.22 MiB)  # without, thus Zygote.pullback
-
-julia> @btime gradient(x -> sum(abs2, (identity∘cbrt).(x)), $xs);
-  55.290 ms (830060 allocations: 49.75 MiB)  # slow path
-  14.747 ms (240043 allocations: 7.25 MiB)   # with `map` rule as before -- better!
-
-# Compare unfused
-
-julia> @btime gradient(x -> sum(abs2, identity.(cbrt.(x))), $xs);
-  69.458 μs (50 allocations: 392.09 KiB)  # fast path -- two copies forward, two back
-  75.041 μs (46 allocations: 470.11 KiB)  # slow path -- 5 copies
-  135.541 μs (27 allocations: 1.30 MiB)   # with `map` rule as before -- worse
-
-
-# Lazy +,-,* for partial fusing
-
-julia> @btime Zygote.gradient(x -> sum(abs2, exp.(2 .* x .- 100)), $xs);
-  81.250 μs (21 allocations: 625.47 KiB)  # special rules + dual numbers, 4 more copies than best
-
-julia> @btime gradient(x -> sum(abs2, exp.(2 .* x .- 100)), $xs);
-  57.166 μs (49 allocations: 470.22 KiB)  # broadcast in *, -
-  54.583 μs (46 allocations: 314.06 KiB)  # broadcasted -- two less copies
-  72.958 μs (26 allocations: 1016.38 KiB) # with `map` rule as before
-
-julia> gradient((x,y) -> sum(abs2, exp.(2 .* x .+ y)), xs, (rand(10)'))
-ERROR: MethodError: no method matching size(::Base.Broadcast.Broadcasted # hmm
-
-julia> @btime gradient((x,y) -> sum(abs2, exp.(x .+ y)), $xs, $(rand(100)'));
-  7.127 ms (75 allocations: 22.97 MiB)   # after
-  12.956 ms (57 allocations: 76.37 MiB)  # before
-
-ulia> @btime Zygote.gradient((x,y) -> sum(abs2, exp.(x .+ y)), $xs, $(rand(100)'));
-  9.937 ms (48 allocations: 45.86 MiB)
-
-=#
+# For certain cheap operations we can easily allow fused broadcast:
 
-# The below is from Zygote: TODO: DO we want to do something better here?
-
-accum_sum(xs::Nothing; dims = :) = NoTangent()
-accum_sum(xs::AbstractArray{Nothing}; dims = :) = NoTangent()
-accum_sum(xs::AbstractArray{<:Number}; dims = :) = sum(xs, dims = dims)
-accum_sum(xs::AbstractArray{<:AbstractArray{<:Number}}; dims = :) = sum(xs, dims = dims)
-accum_sum(xs::Number; dims = :) = xs
-
-# https://github.com/FluxML/Zygote.jl/issues/594
-function Base.reducedim_init(::typeof(identity), ::typeof(accum), A::AbstractArray, region)
-  Base.reducedim_initarray(A, region, NoTangent(), Union{Nothing,eltype(A)})
-end
-
-trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val(ndims(x))))
-
-unbroadcast(x::Union{AbstractArray, Base.Broadcast.Broadcasted}, x̄) =
-  size(x) == size(x̄) ? x̄ :
-  length(x) == length(x̄) ? trim(x, x̄) :
-    trim(x, accum_sum(x̄, dims = ntuple(i -> size(x, i) == 1 ? i : ndims(x̄)+1, Val(ndims(x̄)))))
-
-unbroadcast(x::Number, x̄) = accum_sum(x̄)
-unbroadcast(x::Tuple{<:Any}, x̄) = (accum_sum(x̄),)
-unbroadcast(x::Base.RefValue, x̄) = (x=accum_sum(x̄),)
-
-unbroadcast(x::AbstractArray, x̄::Nothing) = NoTangent()
-
-const Numeric = Union{Number, AbstractArray{<:Number, N} where N}
-
-# function ChainRulesCore.rrule(::typeof(broadcasted), ::typeof(+), xs::Numeric...)
-#     broadcast(+, xs...), ȳ -> (NoTangent(), NoTangent(), map(x -> unbroadcast(x, unthunk(ȳ)), xs)...)
-# end
-
-# Replace Zygote-like fully split broadcasting with one fused over easy operations
 (::∂⃖{1})(::typeof(broadcasted), ::typeof(+), args...) = split_bc_plus(args...)
 (::∂⃖{1})(::typeof(broadcasted), ::typeof(+), arg::Array) = split_bc_plus(arg) # ambiguity
 function split_bc_plus(xs...) where {F}
     broadcasted(+, xs...), Δ -> let Δun = unthunk(Δ)
-        # println("+")
+        _print("broadcast +")
         (NoTangent(), NoTangent(), map(x -> unbroadcast(x, Δun), xs)...)
     end
 end
@@ -188,28 +97,61 @@ Base.eltype(bc::Broadcast.Broadcasted{<:Any, <:Any, typeof(+), <:Tuple}) =
 
 (::∂⃖{1})(::typeof(copy), bc::Broadcast.Broadcasted) = copy(bc), Δ -> (NoTangent(), Δ)
 
-# ChainRulesCore.rrule(::typeof(broadcasted), ::typeof(-), x::Numeric, y::Numeric) = x .- y,
-#   Δ -> let Δ=unthunk(Δ); (NoTangent(), NoTangent(), unbroadcast(x, Δ), -unbroadcast(y, Δ)); end
-
 function (::∂⃖{1})(::typeof(broadcasted), ::typeof(-), x, y)
     broadcasted(-, x, y), Δ -> let Δun = unthunk(Δ)
-        # println("-")
+        _print("broadcast -")
         (NoTangent(), NoTangent(), unbroadcast(x, Δun), -unbroadcast(y, Δun))
         # Ideally you could fuse the - into unbroadcast, mapreduce() not sum, when y is a smaller array
     end
 end
 
-# ChainRulesCore.rrule(::typeof(broadcasted), ::typeof(*), x::Numeric, y::Numeric) = x.*y,
-#   z̄ -> let z̄=unthunk(z̄); (NoTangent(), NoTangent(), unbroadcast(x, z̄ .* conj.(y)), unbroadcast(y, z̄ .* conj.(x))); end
-
 using LinearAlgebra: dot
 
 function (::∂⃖{1})(::typeof(broadcasted), ::typeof(*), x, y)
     broadcasted(*, x, y), Δ -> let Δun = unthunk(Δ)
-        # println("*")
-        dx = x isa Number ? dot(y, Δun) : unbroadcast(x, Δun .* conj.(y))
-        dy = y isa Number ? dot(x, Δun) : unbroadcast(y, Δun .* conj.(x))
+        _print("broadcast *")
+        dx = eltype(x)==Bool ? NoTangent() : x isa Number ? dot(y, Δun) : unbroadcast(x, Δun .* conj.(y))
+        dy = eltype(y)==Bool ? NoTangent() : y isa Number ? dot(x, Δun) : unbroadcast(y, Δun .* conj.(x))
         # When x is an array but a smaller one, instead of dot you may be able to use mapreduce()
+        # Will things like this work? Ref([1,2]) .* [1,2,3]
         (NoTangent(), NoTangent(), dx, dy)
     end
 end
+
+(::∂⃖{1})(::typeof(broadcasted), ::typeof(conj), x) =
+    broadcasted(conj, x), Δ -> (NoTangent(), conj(unthunk(Δ)))
+(::∂⃖{1})(::typeof(broadcasted), ::typeof(conj), x::AbstractArray{Real}) =
+    x, Δ -> (NoTangent(), Δ)
+
+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)
+        ProjectTo(x)(sum(dx; dims))
+    end
+end
+unbroadcast(x::Base.AbstractArrayOrBroadcasted, dx::NoTangent) = NoTangent()
+
+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
+    else
+        sum(dx; dims=2:ndims(dx))
+    end
+    ProjectTo(x)(NTuple{length(x)}(val)) # Tangent
+end

test/runtests.jl

@@ -215,13 +215,36 @@ z45, delta45 = frule_via_ad(DiffractorRuleConfig(), (0,1), x -> log(exp(x)), 2)
 @test delta45 ≈ 1.0
 
 # Broadcasting
-@test gradient(x -> sum(x ./ x), [1,2,3]) == ([1,1,1],)
-@test gradient(x -> sum(exp.(log.(x))), [1,2,3]) == ([1,1,1],)  # derivatives_given_output
-@test gradient(x -> sum((exp∘log).(x)), [1,2,3]) == ([1,1,1],)  # stores pullback
+@test gradient(x -> sum(x ./ x), [1,2,3]) == ([0,0,0],)  # derivatives_given_output
+@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
+exp_log(x) = exp(log(x))
+@test gradient(x -> sum(exp_log.(x)), [1,2,3]) == ([1,1,1],) 
 @test gradient((x,y) -> sum(x ./ y), [1 2; 3 4], [1,2]) == ([1 1; 0.5 0.5], [-3, -1.75])
 @test gradient((x,y) -> sum(x ./ y), [1 2; 3 4], 5) == ([0.2 0.2; 0.2 0.2], -0.4)
-@test gradient(x -> sum(x .> 2), [1,2,3]) == (ZeroTangent(),)   # Bool shortcut
+@test gradient(x -> sum((y -> y/x).([1,2,3])), 4) == (-0.375,)  # closure
+
+@test gradient(x -> sum(sum, (x,) ./ x), [1,2,3])[1] ≈ [-4.1666, 0.3333, 1.1666] atol=1e-3  # array of arrays
+@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 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())
+@test gradient(x -> sum(x .+ [1,2,3]), true) == (NoTangent(),)  # Bool input
+@test gradient(x -> sum(x ./ [1,2,3]), [true false]) == (NoTangent(),)
+@test gradient(x -> sum(x .* [1,2,3]'), (true, false)) == (NoTangent(),)
+
+tup_adj = gradient((x,y) -> sum(2 .* x .+ log.(y)), (1,2), [3,4,5]')
+@test tup_adj[1] == Tangent{Tuple{Int64, Int64}}(6.0, 6.0)
+@test tup_adj[2] ≈ [0.6666666666666666 0.5 0.4]
+@test tup_adj[2] isa Adjoint
+@test gradient(x -> sum(atan.(x, (1,2,3))), Diagonal([4,5,6]))[1] isa Diagonal
+
+@test_broken gradient(x -> sum(gradient(x -> sum(exp.(x)), x)[1]), [1,2,3])  # path 0, MethodError: no method matching Diffractor.Jet(::Int64, ::Float64, ::Tuple{Float64, Float64})
+@test_broken gradient(x -> sum(gradient(x -> sum(x' .* x), x)[1]), [1,2,3]) == ([6,6,6],) # Control flow support not fully implemented yet for higher-order reverse mode
+@test_broken gradient(x -> sum(gradient(x -> sum(x' ./ x.^2), x)[1]), [1,2,3])[1] ≈ [27.675925925925927, -0.824074074074074, -2.1018518518518516]
 
 # Higher order control flow not yet supported (https://github.com/JuliaDiff/Diffractor.jl/issues/24)
 #include("pinn.jl")