add Embedding layer

[CarloLucibello]

Basic implementation.

Maybe could be improved when FluxML/NNlib.jl#255 lands

PR Checklist

• Tests are added
• Entry in NEWS.md
• Documentation, if applicable
• Final review from @dhairyagandhi96 (for API changes).

[DhairyaLGandhi]

Needs GPU tests

[DhairyaLGandhi]

For such a straightforward routine, do we need a layer? I guess people are used to seeing them elsewhere, but why should that make it a necessary thing.

Might just make it a necessary evil. Oh well.

[darsnack]

Nice work

[darsnack]

Seems simple, makes sense to add.

[DhairyaLGandhi]

Yeah, I don't think we want to take the charge on mat here. That, the hard typing for the same logic, and changes to dense aside, it's looking neat.

[CarloLucibello]

Since I need to define and use mat here for the Embedding layer it is needlessly painful to factor out the change to the dense layer, so please let me keep that change to dense here. We don't need to document and export mat

[darsnack]

mat seems useful for writing N-dimensional behavior for layers. I am fine leaving as is and undocumented/unexported. If there is protest, let's just rename it to something else (it's basically a sibling to flatten). Could just have a kwarg in flatten.

[DhairyaLGandhi]

please take care of the rest of the comments in the meantime?

I don't see how mat was something needed to be defined. The reshape api is pretty clever so as to be self documenting, but I can deal with that later.

[CarloLucibello]

please take care of the rest of the comments in the meantime?

don't see anything to take care of

I don't see how mat was something needed to be defined. The reshape api is pretty clever so as to be self documenting, but I can deal with that later.

it is defined because it is convenient, same as flatten, when a code pattern is used in multiple spots it is typically factorized into a function, this is how people write programs

[DhairyaLGandhi]

hard typing for the same logic, and changes to dense aside

These ones.

[CarloLucibello]

hard typing for the same logic,

I don't know what you refer to here

and changes to dense aside

I asked to keep the change here

[CarloLucibello]

@mcabbott given the aliasing problem, I think that the thing to do here is to restrict the input to cpu only (i.e. error out on CuArrays) for the time being so that we can merge this safely. When we have the gather/scatter infrastructure in place in NNlib.jl and CUDA.jl we switch the backend here.

[jeremiedb]

Is there a reason to favor the reliance on ScatterNNLib for the embedding operations rather than having a fix for the incorrect gradient when indices have repeats?
I understand that ScaterNNLib presents the benefit of doing the correct calculation, but it seems to leave in place a nasty silent bug. Having improper gradient if indexing an array with repeating indices looks like the elephant in room.

It's my understanding that JuliaLang/julia#31407 was to address it, but it looks stalled. Any idea how to raise awareness on this issue would be welcome!

[CarloLucibello]

Yes, it would be good to fix FluxML/Zygote.jl#821.
Since the cpu gradient is working correctly, maybe it is enough to address JuliaGPU/CUDA.jl#89

[CarloLucibello]

this can now be implemented on top of the scatter/gather functions implemented in NNlib

[darsnack]

This layer needs to special case (like the normalization layers) in Flux.outputsize since indexing with Nil is non-sensical. We ran into this issue trying to use the layer in FastAI.jl.

Something like

(m::Embedding)(x::AbstractVector{<:Nil}) = fill(nil, size(m.weight, 1), length(x))
(m::Embedding)(x::AbstractArray{<:Nil}) = fill(nil, size(m.weight, 1), last(size(x)))

[CarloLucibello]

gpu failure seems unrelated

[CarloLucibello]

@DhairyaLGandhi can you dismiss your change request?

[bors]

This PR was included in a batch that successfully built, but then failed to merge into master. It will not be retried.

Additional information:

{"message":"1 review requesting changes and 1 approving review by reviewers with write access.","documentation_url":"https://docs.github.com/articles/about-protected-branches"}

[DhairyaLGandhi]

I've left a couple thoughts.

There seem to be a lot of changes around code not intended to be changed in this PR (around the field names of Dense) which should go in a different PR.

[DhairyaLGandhi]

Also needs GPU tests with one-hot and integers.

[mcabbott]

If the plan is to make this perfect, then #1656 should be closed and all comments there addressed here. Alternatively, if the plan is to merge this as a stepping stone to #1656, as mentioned above, then it's a waste of time trying to make sure it doesn't have a hair out of place.

[DhairyaLGandhi]

My comment pointed out adding tests for the features this PR wants to bring in.

#1656 aims to build on top of this, and has its own opinions about API, not implementation.

[CarloLucibello]

need an approval again

[DhairyaLGandhi]

bors r+

[bors]

Build succeeded:

• buildkite/flux-dot-jl