add @adjoin for relu to increase speed

[findmyway]

using Flux, Zygote, BenchmarkTools

x = rand(128)

# before
@benchmark gradient((x) -> sum(relu.(x)), $x)

#=
BenchmarkTools.Trial: 
  memory estimate:  15.41 KiB
  allocs estimate:  425
  --------------
  minimum time:     9.980 μs (0.00% GC)
  median time:      10.649 μs (0.00% GC)
  mean time:        12.495 μs (4.12% GC)
  maximum time:     2.701 ms (98.24% GC)
  --------------
  samples:          10000
  evals/sample:     1
=#

# after

@benchmark gradient((x) -> sum(relu.(x)), $x)
#=
BenchmarkTools.Trial: 
  memory estimate:  2.30 KiB
  allocs estimate:  3
  --------------
  minimum time:     258.013 ns (0.00% GC)
  median time:      303.099 ns (0.00% GC)
  mean time:        368.242 ns (16.51% GC)
  maximum time:     4.796 μs (85.99% GC)
  --------------
  samples:          10000
  evals/sample:     394
=#

[MikeInnes]

This is useful, but you probably don't need the loop vs. just using another broadcast in the backwards pass. See also the gradients in broadcast.jl for examples. The ifelse function might also be useful.

[mcabbott]

This is neat. While you’re in there, perhaps add something similar for identity, as this is as slow as relu before this PR: @btime gradient(x -> sum(identity.(x)), $x);. (And gets done by Dense etc. by default.)

[pevnak]

I have handwritten gradients for sigmoid, tanh and selu, as I needed those for one of my projects. Would it make sense to add those?

explicitgrad(::typeof(tanh), x) = 1f0 - tanh(x)^2
explicitgrad(::typeof(NNlib.σ), x) = σ(x)*(1f0 - σ(x))
function explicitgrad(::typeof(NNlib.selu), x) 
  λ = oftype(x/1, 1.0507009873554804934193349852946)
  α = oftype(x/1, 1.6732632423543772848170429916717)
  λ * ifelse(x > 0, 1, α * exp(x))
end

[MikeInnes]

Thanks @findmyway!

@pevnak yeah, those would be great to have as a PR.

bors r+

[bors]

Build succeeded

• ci/gitlab/gitlab.com