求Numpy数组索引访问函数f(x,item)的导数df_dx实现方案

我需要为以下Numpy数组索引访问函数推导导数实现：

def f(x, item):
  return x[item]

def df_dx(x, item):
  pass

其中item的类型遵循Numpy文档规范，可为int、list、slice或slice列表。该场景存在诸多边界情况，我认为可以通过张量乘法实现，但自己编写的代码并未生效，尝试的代码如下：

to_length = lambda i: i if type(i) == int else len(
            range(i.start, (i.stop if i.stop is not None else -1), (i.step if i.step is not None else 1)))

input_shape = list(X.shape)
seg_shape = []
if type(self.item) == int or type(self.item) == slice:
    seg_shape = [to_length(self.item)]
elif type(self.item) == list:
    seg_shape = [to_length(i) for i in self.item]
else:
    raise Exception(f"Unknown type {type(self.item)}")

v = np.zeros(seg_shape + input_shape)
if len(seg_shape) == 1:
    # np.fill_diagonal(v[:, self.i:self.i + self.n], 1)
    np.fill_diagonal(v[..., self.item], 1)
    return v
elif len(seg_shape) == 2:
    np.fill_diagonal(v[:, :, self.item[0], self.item[1]], 1)
    return v

内容的提问来源于stack exchange，提问作者user155370