Добавление / умножение тензоров и векторов по определенной оси

Question

Как следующая функция может быть эффективно реализована в тензорном потоке и / или 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

Просто добавьте размеры длины единицы 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;Вы можете рассмотреть уменьшение k на 1 и назвать его «осью».

Answer 2

Вы можете автоматически сгенерировать v[:,None,None] для любого измерения Т и любой оси 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.]]])