merge into NNlib and CUDA?

[CarloLucibello]

Hi,
in FluxML/Flux.jl#1431 there was some talk about having the primitives defined here more widely available in the ecosystem. In order to do this, the Zygote and CUDA dependencies should be dropped, because they could be an unnecessary and huge payload for other packages. Therefore, we should have the following steps:

1. Replace Zygote and ZygoteRules adjoint's definitions with ChainRules ones
2. Move the cpu implementations to NNlib.jl
3. Move the gpu kernels to CUDA.jl (which already depends on NNlib.jl) if @maleadt is willing to accept them

@yuehhua does this plan make sense?

cc @dfdx @jeremiedb @chengchingwen

[dfdx]

Currently, NNlib doesn't depend on ChainRules, so it may be better to move adjoints definitions to Flux directly.
I also suggest splitting pullbacks into forward and reverse functions so that they could be used separately by libraries with other gradient calculation rules (e.g. Yota). For example instead of this:

@adjoint function scatter_add!(ys::AbstractArray, us::AbstractArray, xs::AbstractArray)
    ys_ = copy(ys)
    scatter_add!(ys_, us, xs)
    ys_, Δ -> (Δ, gather(Δ, xs), nothing)
end

have this (assuming I understood semantics of @adjoint correctly):

∇scatter_add_ys!(Δ, xs) = Δ
∇scatter_add_us!(Δ, xs) = gather(Δ, xs)

function rrule(::typeof(scatter_add!), args...; kwargs...)
    ys_ = copy(ys)
    scatter_add!(ys_, us, xs)
    ys_, Δ -> (∇scatter_add_ys!(Δ, xs), ∇scatter_add_us!(Δ, xs), nothing)
end

If this looks good to everyone, I can try it out in NNlib / CUDA / Flux during this or next weekend.

[yuehhua]

It makes sense to me. If there is anything I can help, just let me know.
Currently, as @dfdx suggested, put adjoints definitions to Flux directly and ChainRules's definitions to NNlib and CUDA separately?
Also, separate forward and reverse functions?
I will generalize scatter operations to every dimensions as well.

[chengchingwen]

@dfdx Some gradient function need the intermediate values from the forward pass. If we are going to split the definition, then some backward function would need extra argument to get the values instead of recalculate again.

[CarloLucibello]

The only correction to the comments above is that NNlib's rules are currently being moved to NNlib itself FluxML/NNlib.jl#242, so scatter's rules should go there as well.

@yuehhua it would be nice if you could file the PR to NNlib yourself so that you preserve authorship.

[dfdx]

@dfdx Some gradient function need the intermediate values from the forward pass. If we are going to split the definition, then some backward function would need extra argument to get the values instead of recalculate again.

True, and in general some refactoring of forward pass functions may be needed. But in ScatterNNlib all adjoint definitions follow the same, simple to split pattern.

All in all, it looks much easier to have separate forward and reverse pass functions and combine them in a pullback than to have only pullback and try to extract forward & reverse passes from it. This is essentially the reason Yota.jl (and perhaps any non-pullback-based library) still doesn't use ChainRules.jl.

(I hope it doesn't sound like a selfish argument :))

[yuehhua]

Record current status:

• scatter
  • cpu, forward: Add scatter operations FluxML/NNlib.jl#255
  • cpu, backward: Add rrule for scatter FluxML/NNlib.jl#297, scatter cleanup and fix min/max FluxML/NNlib.jl#319
  • cuda, forward: Add scatter for CUDA support FluxML/NNlib.jl#296, Add scatter for CUDA support FluxML/NNlibCUDA.jl#1
  • cuda, backward: Support scatter for CUDA gradient FluxML/NNlibCUDA.jl#13
• gather
  • cpu, forward: add gather FluxML/NNlib.jl#280
  • cpu, backward: Add gradient for gather FluxML/NNlib.jl#318
  • cuda, forward: Gather for CUDA support FluxML/NNlibCUDA.jl#8
  • cuda, backward: Support gather for cuda gradient FluxML/NNlibCUDA.jl#12
• Support CartesianIndex: scatter and gather support element type of idx to be CartesianIndex FluxML/NNlib.jl#308

[yuehhua]

All migrations are complete! Thank you everyone.

[CarloLucibello]

Amazing and relentless work, thanks @yuehhua !