Merge pull request #1018 from mcabbott/broadcast_tidy

Tidy up a few things in broadcasting

Author: mcabbott

src/lib/array.jl

@@ -8,7 +8,7 @@ using Distributed: pmap, AbstractWorkerPool
 @adjoint Array(xs::AbstractArray) = Array(xs), ȳ -> (ȳ,)
 @adjoint Array(xs::Array) = Array(xs), ȳ -> (ȳ,)
 
-@nograd ones, zeros, Base.OneTo, Colon(), one, zero, sizehint!
+@nograd ones, zeros, Base.OneTo, Colon(), one, zero, sizehint!, count
 
 @adjoint Base.vect(xs...) = Base.vect(xs...), Δ -> (Δ...,)
 

src/lib/broadcast.jl

@@ -46,7 +46,7 @@ function Base.reducedim_init(::typeof(identity), ::typeof(accum), A::AbstractArr
 end
 
 trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val(ndims(x))))
-trim(x::Tuple, Δ) = ntuple(k -> Δ[k], length(x))
+trim(x::Tuple, Δ) = NTuple{length(x)}(Δ)
 
 unbroadcast(x::AbstractArray, x̄) =
   size(x) == size(x̄) ? x̄ :
@@ -75,23 +75,17 @@ unbroadcast(x::AbstractArray, x̄::Nothing) = nothing
 
 @adjoint broadcasted(::typeof(*), x::Numeric, y::Numeric) = x.*y,
    Δ -> (nothing, unbroadcast(x, Δ .* conj.(y)), unbroadcast(y, Δ .* conj.(x)))
-@adjoint function broadcasted(::typeof(*), x::Number, y::AbstractArray{<:Number})
-  z, back = pullback(*, x, y)  # this uses dot(y,Δ) instead of Δ .* conj.(y)
-  z, Δ -> (nothing, back(Δ)...)
-end
-@adjoint function broadcasted(::typeof(*), x::AbstractArray{<:Number}, y::Number)
-  z, back = pullback(*, x, y)
-  z, Δ -> (nothing, back(Δ)...)
-end
+@adjoint broadcasted(::typeof(*), x::Number, y::AbstractArray{<:Number}) =
+  _pullback(*, x, y)  # this uses dot(y,Δ) instead of sum(Δ .* conj.(y))
+@adjoint broadcasted(::typeof(*), x::AbstractArray{<:Number}, y::Number) = 
+  _pullback(*, x, y)
 
 @adjoint function broadcasted(::typeof(/), x::Numeric, y::Numeric)
   res = x ./ y
   res, Δ -> (nothing, unbroadcast(x, Δ ./ conj.(y)), unbroadcast(y, .-Δ .* conj.(res ./ y)))
 end
-@adjoint function broadcasted(::typeof(/), x::AbstractArray{<:Number}, y::Number)
-  z, back = pullback(/, x, y)
-  z, Δ -> (nothing, back(Δ)...)
-end
+@adjoint broadcasted(::typeof(/), x::AbstractArray{<:Number}, y::Number) =
+  _pullback(/, x, y)
 
 @adjoint function broadcasted(::typeof(Base.literal_pow), ::typeof(^), x::Numeric, exp::Val{p}) where p
   y = Base.literal_pow.(^, x, exp)
@@ -106,10 +100,10 @@ end
 end
 
 @adjoint broadcasted(::typeof(conj), x::Numeric) =
-  conj.(x), z̄ -> (nothing, conj.(z̄))
+  conj(x), z̄ -> (nothing, conj(z̄))
 
 @adjoint broadcasted(::typeof(real), x::Numeric) =
-  real.(x), z̄ -> (nothing, real.(z̄))
+  real(x), z̄ -> (nothing, real(z̄))
 
 @adjoint broadcasted(::typeof(imag), x::Numeric) =
   imag.(x), z̄ -> (nothing, im .* real.(z̄))
@@ -180,10 +174,9 @@ _dual_safearg(x) = false
   T = Broadcast.combine_eltypes(f, args)
   # Avoid generic broadcasting in two easy cases:
   if T == Bool
-    return f.(args...), _->nothing 
+    return (f.(args...), _ -> nothing) 
   elseif T <: Real && isconcretetype(T) && _dual_purefun(F) && all(_dual_safearg, args)
-    y, back = broadcast_forward(f, args...)
-    return y, ȳ -> (nothing, nothing, back(ȳ)...)
+    return broadcast_forward(f, args...)
   end
   len = inclen(args)
   y∂b = _broadcast((x...) -> _pullback(__context__, f, x...), args...)
@@ -195,7 +188,7 @@ _dual_safearg(x) = false
     end
     (nothing, accum_sum(dxs[1]), map(unbroadcast, args, Base.tail(dxs))...)
   end
-  y, ∇broadcasted
+  return y, ∇broadcasted
 end
 
 @adjoint function broadcasted(::AbstractArrayStyle{0}, f, args...)
@@ -226,6 +219,7 @@ using ForwardDiff: Dual
 
 dual(x, p) = x
 dual(x::Real, p) = Dual(x, p)
+dual(x::Bool, p) = x
 
 function dual_function(f::F) where F
   function (args::Vararg{Any,N}) where N
@@ -237,28 +231,36 @@ function dual_function(f::F) where F
 end
 
 @inline function broadcast_forward(f, args::Vararg{Any,N}) where N
-  T = Broadcast.combine_eltypes(f, args)
+  valN = Val(N)
   out = dual_function(f).(args...)
   eltype(out) <: Dual || return (out, _ -> nothing)
   y = map(x -> x.value, out)
-  _back(ȳ, i) = unbroadcast(args[i], ((a, b) -> a*b.partials[i]).(ȳ, out))
-  back(ȳ) = ntuple(i -> _back(ȳ, i), N)
-  return y, back
+  function bc_fwd_back(ȳ)
+    dargs = ntuple(valN) do i
+      unbroadcast(args[i], broadcast((y1, o1) -> y1 * o1.partials[i], ȳ, out))
+    end
+    (nothing, nothing, dargs...) # nothings for broadcasted & f
+  end
+  return y, bc_fwd_back
 end
 
 @init @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" begin
+
   const CuArrayStyle = CUDA.AbstractGPUArrayStyle
 
-  if isdefined(CUDA, :cufunc)
-    @eval @adjoint function broadcasted(::CuArrayStyle, f, args...)
-      y, back = broadcast_forward(CUDA.cufunc(f), args...)
-      y, ȳ -> (nothing, nothing, back(ȳ)...)
-    end
-  else # CUDA >= 3.0
-    @eval @adjoint function broadcasted(::CuArrayStyle, f, args...)
-      y, back = broadcast_forward(f, args...)
-      y, ȳ -> (nothing, nothing, back(ȳ)...)
-    end
+  if isdefined(CUDA, :cufunc)  # CUDA < 3.0
+
+    @eval @adjoint broadcasted(::CuArrayStyle, f, args...) =
+      broadcast_forward(CUDA.cufunc(f), args...)
+
+  else # CUDA >= 3.0 -- don't need cufunc(f). 
+       # Ordinary broadcasting calls broadcast_forward anyway when certain its' safe,
+       # so perhaps this can be deleted? Possible edge case here:
+       # https://github.com/FluxML/Zygote.jl/pull/1018#issuecomment-873629415
+
+    @eval @adjoint broadcasted(::CuArrayStyle, f, args...) =
+      broadcast_forward(f, args...)
+
   end
 
   @adjoint CUDA.CuArray{N,T}(xs::Array) where {N,T} =

test/cuda.jl

@@ -9,7 +9,7 @@ CUDA.allowscalar(false)
   @test gradient(x -> sum(cu(x)), r)[1] isa Array{Float32, 2}
 end
 
-@testset "basic bcasting" begin
+@testset "broadcasting" begin
   a = Float32.(1:9)
   a_gpu = a |> cu
 
@@ -24,6 +24,13 @@ end
   g_gpu = gradient(x -> w(x), a_gpu)[1]
   @test g_gpu isa CuArray
   @test g_gpu |> collect ≈ g
+
+  # https://github.com/FluxML/Zygote.jl/issues/1027                  # status on Zygote v0.6.14, CUDA v3.3.0 in comments:
+  @test gradient(x -> sum(x .!= 0), a_gpu) == (nothing,)             # was MethodError: no method matching iterate(::Nothing)
+  @test gradient(x -> sum(x .> 3), a_gpu) == (nothing,)
+  g3 = gradient(x -> sum(x .^ 3) / count(x .> 3), a)[1]              # was Can't differentiate gc_preserve_end expression
+  @test_skip cu(g3) ≈ gradient(x -> sum(x .^ 3) / sum(x .> 3), a_gpu)[1]  # was KernelException -- not fixed by PR #1018
+  @test cu(g3) ≈ gradient(x -> sum(x .^ 3) / count(x .> 3), a_gpu)[1] 
 end
 
 @testset "sum(f, x)" begin