# Create an array (a) of shape 3, 4, 8 (K=3, J=4, I=8). tidx is an array of the same length as a.shape[1], i.e. contains J = 4 elements where each index denotes which element of K should be chosen. 
 import numpy as np
a = np.random.randint(0, 10, (3, 4, 8))
print("Original array and shape:")
print(a)
print(a.shape)
print("--------------------------------")
tidx = np.random.randint(0, 3, 4)
print("tidex: ",tidx)
print("Result:")
print(a[tidx, np.arange(len(tidx)),:])
