Перемещение строк символической матрицы Теано

Question

Я выполняю параллельный отбор образцов Гиббса с использованием Theano. Я пытаюсь создать функцию Anano, которая принимает матрицу X и меняет некоторые ее строки. У меня есть символический двоичный вектор с именем swaps, который обозначает, какие строки должны быть заменены (то есть, если swaps[i] == 1, то X[i] и X[i+1] следует поменять местами). Порядок обмена не важен для меня.

Я пытался написать theano.scan, который проходит через вектор swaps и выполняет обмен X строк за строкой. Проблема в том, что Theano не позволяет делать что-то вроде X[pos], X[pos + 1] = X[pos + 1], X[pos] с символическими переменными. Вот простой фрагмент кода того, что я пытаюсь сделать.

import numpy as np 

import theano 
import theano.tensor as T 

def swap(swp, pos, idx): 
    if swp: idx[pos], idx[pos + 1] = idx[pos + 1], idx[pos] 
    return idx 

max_length = 10 
swaps = T.ivector('swaps') 
idx = T.ivector('idx') 
pos = T.iscalar('pos') 

new_idx, updates = theano.scan(swap, 
           sequences=[swaps, T.arange(max_length)], 
           outputs_info=idx) 

do_swaps = theano.function([swaps, idx], new_idx[-1], updates=updates) 

idx_swapped = do_swaps(np.array([1, 1, 0, 1]), np.arange(5)) 
print idx_swapped 

Любые идеи о том, как я могу сделать это правильно?

Answer 1

Хорошо, вот очень простое решение, которое я нашел.

import numpy as np 

import theano 
import theano.tensor as T 

def swap(swp, pos, X): 
    return T.concatenate([X[:pos],X[[pos+swp]],X[[pos+1-swp]],X[pos+2:]]) 

max_length = 10 
swaps = T.ivector('swaps') 
pos = T.iscalar('pos') 
X  = T.vector('X') 

new_X, _ = theano.scan(swap, 
         sequences=[swaps, T.arange(max_length)], 
         outputs_info=X) 

do_swaps = theano.function([swaps, X], new_X[-1]) 

X_swapped = do_swaps(np.array([1, 1, 0, 1], dtype='int32'), np.arange(5)) 
print X_swapped 

Однако я не уверен, насколько это оптимально или нет для выполнения на графическом процессоре.