沿特定轴添加/乘以张量和向量

Question

如何在 tensorflow 和/或 numpy 中有效地实现以下功能？add_along_axis(tensor=T, vector=v, axis=k)     # T is a tensor of shape (N1,...,Nd) (unknown beforehand)     # v is a vector with N components     # k is an integer such that Nk=N     S = T+v, summed along k     return S那是带有分量S的(N1,..,Nd)张量S[i1,...,id]=T[i1,...,id] + v[ik]请注意，任何数量的Nj,j≠k可能巧合地等于N，因此标准广播不是一种选择。例如：让T = np.zeros( (3,3,3) )和v = [1,2,3]那么正确的输出应该是f(T,v,1) = [[[1., 1., 1.],  [[2., 2., 2.],  [[3., 3., 3.],             [1., 1., 1.],   [2., 2., 2.],   [3., 3., 3.],             [1., 1., 1.]],  [2., 2., 2.]],  [3., 3., 3.]]] f(T,v,2) = [[[1., 1., 1.],  [[1., 1., 1.],  [[1., 1., 1.],             [2., 2., 2.],   [2., 2., 2.],   [2., 2., 2.],             [3., 3., 3.]],  [3., 3., 3.]],  [3., 3., 3.]]]f(T,v,3) = [[[1., 2., 3.],  [[1., 2., 3.],  [[1., 2., 3.],             [1., 2., 3.],   [1., 2., 3.],   [1., 2., 3.],             [1., 2., 3.]],  [1., 2., 3.]],  [1., 2., 3.]]]在这里，目标行为可以分别通过编写T+v[:,None,None]、T+v[None,:,None]和来实现T+v[None,None,:]。但是，我不知道这种方法在张量形状未预定义的情况下如何工作。

Answer 1

您可以通过执行列表理解v[:,None,None]为 T 的任何维度和任何轴自动生成 ，k例如：

def f(T,v,k):

    return T+v[[np.newaxis if i+1 != k else slice(None) for i in range(T.ndim) ]]

np.newaxis等价于None，slice(None)等价于:。结果如预期：

print (f(T,v,2))

array([[[1., 1., 1.],

        [2., 2., 2.],

        [3., 3., 3.]],

       [[1., 1., 1.],

        [2., 2., 2.],

        [3., 3., 3.]],

       [[1., 1., 1.],

        [2., 2., 2.],

        [3., 3., 3.]]])

Answer 2

只需将 T.ndim-k 单位长度维度附加到 v 和 numpy 的广播规则按设计工作：

def f(T, v, k):

    v = asarray(v)

    return T + v.reshape(v.shape + (1,)*(T.ndim-k))

请注意，您对 k 的定义比标准 numpy 轴编号大 1；您可能会考虑将 k 减少 1 并将其称为“轴”。