基于偏移量向量的tensor3元素位置偏移

Question

我有一个tensor3 (即三维数组) x。

[[[ 0  1  2  3]
  [ 4  5  6  7]
  [ 8  9 10 11]]

 [[12 13 14 15]
  [16 17 18 19]
  [20 21 22 23]]]

以及y向量(即一维数组)，我们将其称为“偏移量”向量，因为它指定了所需的偏移量：

[2, 1]

我希望基于向量x调整y元素的位置，以便输出如下(在第二维度上执行移位)：

[[[ a  b  c  d]
  [ e  f  g  h]
  [ 0  1  2  3]]

 [[ i  j  k  l]
  [12 13 14 15]
  [16 17 18 19]]]

a，b，…，l可能是任意数字。

例如，一个有效的输出可以是：

[[[ 0  0  0  0]
  [ 0  0  0  0]
  [ 0  1  2  3]]

 [[ 0  0  0  0]
  [12 13 14 15]
  [16 17 18 19]]]

另一个有效的输出可以是：

[[[ 4  5  6  7]
  [ 8  9 10 11]
  [ 0  1  2  3]]

 [[20 21 22 23]
  [12 13 14 15]
  [16 17 18 19]]]

我知道函数theano.tensor.roll(x, shift, axis=None)，但是shift只能接受一个标量作为输入，即它移动所有具有相同偏移量的元素。

例如，代码：

import theano.tensor
from theano import shared
import numpy as np

x = shared(np.arange(24).reshape((2,3,4)))
print('theano.tensor.roll(x, 2, axis=1).eval(): \n{0}'.
      format(theano.tensor.roll(x, 2, axis=1).eval()))

产出：

theano.tensor.roll(x, 2, axis=1).eval():
[[[ 4  5  6  7]
  [ 8  9 10 11]
  [ 0  1  2  3]]

 [[16 17 18 19]
  [20 21 22 23]
  [12 13 14 15]]]

这不是我想要的。

如何根据偏移量向量移动tensor3元素的位置？(注意，在本例中提供的代码中，为了方便起见，tensor3是一个共享变量，但在我的实际代码中，它将是一个符号变量)

Answer 1

为此，我找不到任何专用函数，所以我只使用了theano.scan。

import theano
import theano.tensor

from theano import shared
import numpy as np

y = shared(np.array([2,1]))
x = shared(np.arange(24).reshape((2,3,4)))
print('x.eval():\n{0}\n'.format(x.eval()))

def shift_and_reverse_row(matrix, y):    
    '''
    Shift and reverse the matrix in the direction of the first dimension (i.e., rows)
    matrix: matrix 
    y: scalar
    '''
    new_matrix = theano.tensor.zeros_like(matrix)
    new_matrix = theano.tensor.set_subtensor(new_matrix[:y,:], matrix[y-1::-1,:])
    return new_matrix

new_x, updates = theano.scan(shift_and_reverse_row, outputs_info=None,
                             sequences=[x, y[::-1]] )
new_x = new_x[:, ::-1, :]
print('new_x.eval(): \n{0}'.format(new_x.eval()))

产出：

x.eval():
[[[ 0  1  2  3]
  [ 4  5  6  7]
  [ 8  9 10 11]]

 [[12 13 14 15]
  [16 17 18 19]
  [20 21 22 23]]]

new_x.eval():
[[[ 0  0  0  0]
  [ 0  0  0  0]
  [ 0  1  2  3]]

 [[ 0  0  0  0]
  [12 13 14 15]
  [16 17 18 19]]]