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

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 并将其称为“轴”。