theano ~ use an index matrix and embeddings matrix to produce a 3D tensor?

Question

Let's say I have an index matrix:

index_matrix = [[0,1],
                [2,1]]

And an embeddings matrix:

embeddings_matrix = [[1,2,3],
                     [2,3,4],
                     [4,5,6],
                        .
                        .
                        .   ]

Each element in the index_matrix corresponds to a particualr row in the embeddings_matrix. For example, 0 in the index_matrix corresponds to [1,2,3] in the embeddings_matrix.

Using theano syntax only, how can I build a 3D tensor, where each index in the index_matrix is replaced with a row or vector from the embeddings_matrix?

Edit: I was able to solve a similar problem: that of using an index_vector = [0,1,2,1] and the embeddings_matrix above to generate a 2D tensor. For that, I needed the following theano syntax:

Matrix = embeddings_matrix[index_vector].reshape((index_vector.shape[0], -1))

However, I am a bit stuck on the 3D case.

Answer 1

You can simply index into the embedding matrix with the index matrix:

import numpy
import theano
import theano.tensor as tt

embeddings_matrix = theano.shared(numpy.array([[1, 2, 3], [2, 3, 4], [4, 5, 6]], dtype=theano.config.floatX))
index_matrix = tt.imatrix()
y = embeddings_matrix[index_matrix]
f = theano.function([index_matrix], y)
output = f(numpy.array([[0, 1], [2, 1]], dtype=numpy.int32))
print output.shape, '\n', output

This prints

(2L, 2L, 3L) 
[[[ 1.  2.  3.]
  [ 2.  3.  4.]]

 [[ 4.  5.  6.]
  [ 2.  3.  4.]]]

The output is 3D with each element in the index matrix replaced with the corresponding embedding vector as required.