Add EnzymeRule for conv!/gather!/scatter!/dropout!/pool!

[wsmoses]

Trying this out, hopefully to serve as a template for other functions.

[wsmoses]

I wasn't sure where to put the test but here's an example:

using Enzyme
using NNlib

w = randn(Float32, 3, 3, 5, 7);
dw = zero(w)
loss(w, x) = sum(conv(x, w))
x = randn(Float32, (3, 3, 5, 8));

Enzyme.autodiff(Reverse, loss, Duplicated(w, dw), Const(x));

@show dw

[wsmoses]

Another note, right now there is almost certainly more copying of inputs than necessary.

In the future we should add some sort of (type?) annotation to denote that some memory is immutable, or perhaps have a copy-on-write wrapper (as proposed by @gaurav-arya)

[wsmoses]

Sure.

…

On Sun, Sep 24, 2023 at 4:26 AM Carlo Lucibello ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In Project.toml <#536 (comment)>: > @@ -16,13 +17,15 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [weakdeps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" We could depend just on EnzymeCore right? — Reply to this email directly, view it on GitHub <#536 (review)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXE3AG3G2VXUBRUNVR3X374F3ANCNFSM6AAAAAA5EXWQE4> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[CarloLucibello]

Even if EnzimeCore is a lean dependency I think it is better to put all under a NNlibEnzymeCoreExt extension.

[wsmoses]

So I tried that originally, but the problem with that is then your cuda extension would itself need an enzymecore extension.

If it matters, it also mirrors the way chainrulescore is used presently.

[CarloLucibello]

So I tried that originally, but the problem with that is then your cuda extension would itself need an enzymecore extension.

this can be addressed having also a NNlibEnzymeCoreCUDA extension.

If it matters, it also mirrors the way chainrulescore is used presently.

true but chainrulescore is AD framework agnostic and quite widespread, enzymecore is not. It is a very lean dependency, so if making it an extension proves to be too tricky I wouldn't oppose making it a hard one but I think we should avoid it if we can.

[wsmoses]

@CarloLucibello my attempt to use extensions (current code) fails. However, without the extension the previous commit (f430a9b) succeeds.

If you see how to fix the extension, we can use that, if not we can use the prior commit.

[CarloLucibello]

maybe better to factor out this change to another PR

[maxfreu]

Maybe this is a bit early, but I'd love to see benchmarks of this vs Zygote.

[ToucheSir]

This was a tour de force of a PR, thanks @wsmoses!