在theano中使用索引矩阵索引张量?

时间:2015-11-26 22:15:13

标签: python numpy matrix theano

我有一个theano张量A,使得A.shape =(40,20,5)和theano矩阵B,使得B.shape =(40,20)。是否有一个单行操作我可以执行以获得矩阵C,其中C.shape =(40,20)和C(i,j)= A [i,j,B [i,j]]和theano语法?

基本上,我想用B作为索引矩阵;使用theano做什么最有效/优雅?

1 个答案:

答案 0 :(得分:4)

您可以在numpy中执行以下操作:

import numpy as np

A = np.arange(4 * 2 * 5).reshape(4, 2, 5)
B = np.arange(4 * 2).reshape(4, 2) % 5

C = A[np.arange(A.shape[0])[:, np.newaxis], np.arange(A.shape[1]), B]

所以你可以在theano做同样的事情:

import theano
import theano.tensor as T

AA = T.tensor3()
BB = T.imatrix()

CC = AA[T.arange(AA.shape[0]).reshape((-1, 1)), T.arange(AA.shape[1]), BB]

f = theano.function([AA, BB], CC)

f(A.astype(theano.config.floatX), B)