use ProjectTo in broadcasting, etc

Author: mcabbott

src/compiler/chainrules.jl

@@ -128,6 +128,20 @@ Convert `x` from the format Zygote uses internally to differentials types ChainR
   ChainRules.Tangent{Any, typeof(xp)}(xp)
 end
 
+"""
+  _project(x)(dx)
+  _project(x, dx)
+
+The function `_project(x)` returns a projector, which standardises the gradient `dx` for type & shape.
+Uses `ChainRulesCore.ProjectTo`, but is safe to apply to arbitrary input.
+The two-argument `_project(x, dx)` applies this immediately.
+"""
+@inline _project(x) = identity  # fallback: do nothing!
+@inline _project(x::Numeric) = wrap_chainrules_output ∘ ProjectTo(x)
+@inline _project(x::Ref{<:Numeric}) = wrap_chainrules_output ∘ ProjectTo(x)
+
+@inline _project(x, dx) = _project(x)(dx)
+
 """
   ZBack{F}(back) <: Function
 

src/compiler/interface.jl

@@ -73,7 +73,8 @@ julia> gradient([7, 11], 0, 1) do x, y, d
 """
 function gradient(f, args...)
   y, back = pullback(f, args...)
-  return back(sensitivity(y))
+  grad = back(sensitivity(y))
+  map(_project, args, grad)
 end
 
 Base.adjoint(f::Function) = x -> gradient(f, x)[1]
@@ -95,7 +96,8 @@ true
 """
 function withgradient(f, args...)
   y, back = pullback(f, args...)
-  (val = y, grad = back(sensitivity(y)))
+  grad = back(sensitivity(y))
+  (val = y, grad = map(_project, args, grad))
 end
 
 # Param-style wrappers

src/lib/array.jl

@@ -43,7 +43,7 @@ end
     dxv = view(dx, inds...)
     dxv .= accum.(dxv, _droplike(dy, dxv))
   end
-  return (dx, map(_->nothing, inds)...)
+  return (_project(x, dx), map(_->nothing, inds)...)
 end
 
 """

src/lib/broadcast.jl

@@ -45,18 +45,14 @@ function Base.reducedim_init(::typeof(identity), ::typeof(accum), A::AbstractArr
   Base.reducedim_initarray(A, region, nothing, Union{Nothing,eltype(A)})
 end
 
-trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val(ndims(x))))
-trim(x::Tuple, Δ) = NTuple{length(x)}(Δ)
-
 unbroadcast(x::AbstractArray, 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̄)))))
+  length(x) == length(x̄) ? _project(x, x̄) :
+    _project(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::Tuple, x̄) = trim(x, length(x) == length(x̄) ? x̄ : accum_sum(x̄; dims=2:ndims(x̄))) # case length(x) > 1
+unbroadcast(x::Tuple, x̄) =  NTuple{length(x)}(length(x) == length(x̄) ? x̄ : accum_sum(x̄; dims=2:ndims(x̄))) # case length(x) > 1
 
 unbroadcast(x::AbstractArray, x̄::Nothing) = nothing