gpu-friendly

src/tracker/lib.jl

@@ -72,10 +72,14 @@ unbroadcast(x, Δ) =
   size(x) == size(Δ) ? Δ :
     sum(Δ, filter(n -> size(x, n) == 1, 1:ndims(Δ)))
 
-function back!(b::Broadcasted, Δ, args...)
-  Δargs = ntuple(i -> Δ .* getindex.(partials.(b.data), i), length(args))
+function getpartial(Δ, x, i)
+  @inbounds p = getindex(partials(x), i)
+  return Δ * p
+end
+
+function back!(b::Broadcasted, Δ, args::Vararg{Any,N}) where N
+  Δargs = ntuple(i -> getpartial.(Δ, b.data, i), Val{N})
   foreach((x, Δ) -> @back!(x, unbroadcast(x, Δ)), args, Δargs)
-  return
 end
 
 Base.Broadcast._containertype(::Type{<:TrackedArray}) = TrackedArray