在theano中用索引矩阵索引张量？

Question

我有一个 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 做这个最efficient/elegant 是什么？

Answer 1

您可以在 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)

在theano中用索引矩阵索引张量？

Indexing tensor with index matrix in theano?

python

numpy

matrix

theano