support 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):
    # # 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

贡献者指南

技术栈

python

领域

machine learning

议题类型

feature

难度面向新贡献者的预计实现难度，1 表示很小改动，5 表示专家级工作。

3

预计时间有经验贡献者完成调查、实现、测试并准备 pull request 的粗略时间范围。

1-3 hours

活动状态议题当前的可参与程度：新鲜、活跃、陈旧、阻塞或等待维护者输入。

fresh

清晰度议题是否清楚说明期望改动、验收标准和下一步。

clear

前置要求

familiarity with autograd's VJP implementationunderstanding of numpy indexing

新手友好度1-100 的估计分数，表示该议题对首次贡献者的友好程度。

60

研究方向

To implement VJP for numpy.take, first examine how existing VJPs for indexing operations (e.g., getitem ) are defined in autograd's source code, likely in the numpy wrapper module. Since take is a subset of getitem, the VJP can be derived by mapping take calls to equivalent getitem slices. Check the issue comments for any additional guidance. After implementing the VJP, add a test case similar to the provided reproducer to ensure correctness.