Gather

[dfdx]

It has been the hardest and the most time-consuming operation so far, so I really hope this time I've done it correctly 😄

As a side effect, I've added the take() function, which is similar to numpy.take() and seems to be useful e.g. for taking slices of input tensors in NLP models. So we may want to move it to a more general-purpose package such as NNlib once our implementation is battle tested in real graphs.

[ToucheSir]

Sounds like a tour de force! Is there a summary of how take differs from NNlib.gather?

[dfdx]

In short, it's just a different operation.

NNlib.gather can work on individual indices, e.g. you can write:

x = rand(3, 4)
@assert NNlib.gather(y, [(1, 1), (1, 2), (2, 2)]) == [x[1, 1], x[1, 2], x[2, 2]]

For example, it's useful to find implement an NLL loss:

y = softmax(rand(3, 4))   # classifier predictions
c = [1, 3, 1, 2]                 # correct classes
idxs = [(c, i) for (i, c) in enumerate(c)]
loss = mean(NNlib.gather(y, idxs))

NNlib.gather is also the reverse of NNlib.scatter and equivalent of ONNX's GatherND or maybe GatherElements - I haven't figured it out yet.

take() works on slices along a specific axis (dim in Julia version). Say, data is 3D and dim = 2. Then for every entry i in idxs it takes data[:, i, :]. It's like extracting individual slices from a ND array and re-organizing them according to the structure in idxs. ONNX's docs are pretty confusing in this regard, but NumPy explain it well.

take() also reduces to getindex() when data is 1D and length(idxs) == 1 - it turns out this is how ONNX treats expressions like size(X)[1].

I tried t to implement take() using NNlib.gather() at some point, but I spent most of the time figuring out what exactly ONNX's Gather does in each case, and when I got the idea, implementing it using array views turned to be easier than adjusting the indices (but it may still be an option in future).