解开跨步索引

Question

我正在尝试编写一个例程，该例程返回给定两个操作数形状的 raveled 跨步索引。它应该考虑广播，并且从这些“raveled strided indices”中应该可以获取两个数组（操作数）正在访问哪些值（如果有扁平化）。

一些我希望能更好地解释这一点的例子：

shapeA = (4,3)
shapeB = (3,)

应该输出

strideA = [0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11]
strideB = [0,  1,  2,  0,  1,  2,  0,  1,  2,  0,  1,  2]

另一个例子：

shapeA = (4,1)
shapeB = (3,)

应该输出

strideA = [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]
strideB = [0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2]

我的尝试（仅适用于第一个示例）：

def ravel_strided_indices(shapeA, shapeB):
    out_shape = np.broadcast_shapes(shapeA, shapeB)
    flatten = np.prod(out_shape)
    strideA = np.arange(flatten) % np.prod(shapeA)
    strideB = np.arange(flatten) % np.prod(shapeB)
    return strideA, strideB

在 numpy 中是否存在任何方便的方法（可能类似于 np.ravel_multi_index 或 np.indices）？

编辑

我正在寻找可以进行此操作的东西

A = np.random.randn(3,1,5,2,1)
B = np.random.randn(1,4,5,1,5)

A + B

相当于这个操作：

indicesA, indicesB = ravel_strided_indices(A.shape, B.shape)
out_shape = np.broadcast_shapes(A.shape, B.shape)

(A.flatten()[indicesA] + B.flatten()[indicesB]).reshape(out_shape)

Answer 1

也许以下内容可以帮助您入门。

import numpy as np
AA = np.random.randn(3,1,5,2,1)
BB = np.random.randn(1,4,5,1,5)
CC = AA + BB
out_shape = np.broadcast_shapes(AA.shape,BB.shape)
aidx   = np.arange(np.prod(AA.shape)).reshape(AA.shape)
bidx   = np.arange(np.prod(BB.shape)).reshape(BB.shape)
aidx_b = np.broadcast_to(aidx,CC.shape)
bidx_b = np.broadcast_to(bidx,CC.shape)
cnew   = (AA.flatten()[aidx_b.flatten()] + BB.flatten()[bidx_b.flatten()]).reshape(out_shape)
print('norm diff=',np.linalg.norm(CC-cnew))

编辑：当然，问题是你为什么要做这样的事情。