3dim tensor, indexed by 3dim tensor, into 2nd dim

doit.py

import theano.tensor as T
import theano
from theano import shared as S
import numpy as np

'''
meant to be run in an ipython notebook
'''

def f():
  ### simulating:
  #### batch = 3
  #### sequence size = 4
  #### feature size = 5
  A = S(np.arange(60).reshape(3,4,5))
  
  #### here, for each time step, we have 10 things to embed; goal is (3,4,10,5)
  x = S(np.random.randint(0, 4, (3,4,10)))
  ### A.transpose(1,0,2)[x].shape -- > (3, 4, 10, 3, 5)
  ### A.transpose(1,0,2)[x].transpose(0,3,1,2,4).shape ---> (3, 3, 4, 10, 5)
  ### A.transpose(1,0,2)[x].transpose(0,3,1,2,4)[np.arange(A.shape[0]), np.arange(A.shape[0])].shape ---> (3, 4, 10, 5)
  x_emb = A.dimshuffle(1,0,2)[x].dimshuffle(0,3,1,2,4)[T.arange(A.shape[0]), T.arange(A.shape[0])]
  ## numpy version
  #x_emb = A.transpose(1,0,2)[x].transpose(0,3,1,2,4)[np.arange(A.shape[0]), np.arange(A.shape[0])]
%timeit f

def f2():
    A = S(np.arange(60).reshape(3,4,5))
    x = S(np.random.randint(0, 4, (3,4,10)))
    x_emb = A.dimshuffle(1,0,2)[x].diagonal(axis1=0, axis2=3).dimshuffle(3,0,1,2).eval()
%timeit f2