支持 numpy.take

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):
    # # 可以运行:
    # return x[:, idx, :].sum()
    # # 不能运行:
    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` 主要只是一个方便的语法/子集的 `getitem` 索引,所以我想没有什么是根本阻止它的。