use purity check in map, too?

Author: mcabbott

src/lib/array.jl

@@ -215,19 +215,27 @@ _tryreverse(m, x) = x
 _tryreverse(m::typeof(map), x::Union{AbstractVector, Tuple}) = reverse(x)
 
 for (mapfunc,∇mapfunc) in [(:map,:∇map),(:pmap,:∇pmap)]
-  @eval function $∇mapfunc(cx, f, args...)
+  @eval function $∇mapfunc(cx, f::F, args...) where {F}
     ys_and_backs = $mapfunc((args...) -> _pullback(cx, f, args...), args...)
     if isempty(ys_and_backs)
       ys_and_backs, _ -> nothing
     else
-      ys, backs = unzip(ys_and_backs)
+      ys = map(first, ys_and_backs)
       ys, function (Δ)
         isnothing(Δ) && return nothing
-        # Apply pullbacks in reverse order. Needed for correctness if `f` is stateful.
-        Δf_and_args_zipped = $mapfunc((f, δ) -> f(δ), _tryreverse($mapfunc, backs, Δ)...)
-        Δf_and_args = unzip(_tryreverse($mapfunc, Δf_and_args_zipped))
-        Δf = reduce(accum, Δf_and_args[1])
-        (Δf, Δf_and_args[2:end]...)
+        if _purefun(F) && length(args) == 1
+          Δarg = $mapfunc(((_,pb), δ) -> last(pb(δ)), ys_and_backs, Δ) # No unzip needed
+          (nothing, Δarg)
+        elseif _purefun(F)
+          Δargs = unzip($mapfunc(((_,pb), δ) -> Base.tail(pb(δ)), ys_and_backs, Δ))
+          (nothing, Δargs...)
+        else
+          # Apply pullbacks in reverse order. Needed for correctness if `f` is stateful.
+          Δf_and_args_zipped = $mapfunc(((_,pb), δ) -> pb(δ), _tryreverse($mapfunc, ys_and_backs, Δ)...)
+          Δf_and_args = unzip(_tryreverse($mapfunc, Δf_and_args_zipped))
+          Δf = reduce(accum, Δf_and_args[1])
+          (Δf, Δf_and_args[2:end]...)
+        end
       end
     end
   end