Theano расширенная индексация для тензора, общий индекс

Question

У меня есть тензор probs с probs.shape = (max_time, num_batches, num_labels).

И у меня есть тензор targets с targets.shape = (max_seq_len, num_batches), где значения являются индексами меток, то есть для третьего измерения в probs.

Теперь я хочу получить тензор probs_y с probs.shape = (max_time, num_batches, max_seq_len), где третье измерение - это индекс в targets. В основном

probs_y[:,i,:] = probs[:,i,targets[:,i]] 

для всех 0 <= i < num_batches.

Как я могу это достичь?

Аналогичная проблема с решением была опубликована here.

Решение есть, если я правильно понимаю, был бы:

probs_y = probs[:,T.arange(targets.shape[1])[None,:],targets] 

Но это не похоже на работу. Я получаю: IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices.

Кроме того, не является создание временного T.arange немного дорогостоящим? Esp, когда я пытаюсь обходным путем, действительно делая его полным плотным целым массивом. Должен быть лучший способ.

Возможно, theano.map? Но, насколько я понимаю, это не распараллеливает код, так что это тоже не решение.

Answer 1

Это работает для меня:

import theano 
import theano.tensor as T 

max_time, num_batches, num_labels = 3, 4, 6 
max_seq_len = 5 

probs_ = np.arange(max_time * num_batches * num_labels).reshape(
    max_time, num_batches, num_labels) 

targets_ = np.arange(num_batches * max_seq_len).reshape(max_seq_len, 
    num_batches) % (num_batches - 1) # mix stuff up 

probs, targets = map(theano.shared, (probs_, targets_)) 

print probs_ 
print targets_ 

probs_y = probs[:, T.arange(targets.shape[1])[:, np.newaxis], targets.T] 

print probs_y.eval() 

Выше использовал транспонированную версию ваших индексов. Ваше точное предложение также работает

probs_y2 = probs[:, T.arange(targets.shape[1])[np.newaxis, :], targets] 

print probs_y2.eval() 
print (probs_y2.dimshuffle(0, 2, 1) - probs_y).eval() 

Так что, может быть, ваша проблема в другом месте.

Что касается скорости, я затрудняюсь с тем, что может быть быстрее, чем это. map, который является специализацией scan почти наверняка нет. Я не знаю, в какой степени arange фактически построен, а не просто переименован.