support numpy.take

[jcmgray]

Currently tracing the gradient through take is not supported:

import numpy as np
import autograd as ag
import autograd.numpy as anp

rng = np.random.default_rng(42)
x = rng.uniform(size=(3, 4, 5))
idx = rng.integers(0, 4, size=(6,))

def foo(x, idx):
    # # works:
    # return x[:, idx, :].sum()
    # # doesn't work:
    return anp.take(x, idx, axis=1).sum()

gfoo = ag.grad(foo, argnum=0)
gfoo(x, rng.integers(0, 4, size=(6,))).shape
# NotImplementedError: VJP of take wrt argnums (0,) not defined

take is mostly just a convenient syntax/subset of getitem indexing, so I suppose there is nothing fundamentally blocking it.

[fjosw]

Hi @jcmgray, thanks for spotting this. I currently don't have the time to look at this, not sure about the other maintainers, but we are happy about all contributions in case you want to have a closer look yourself.

[agriyakhetarpal]

Indeed, I am happy to review any PRs that you or any other interested contributor might want to put up!