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)),:])

Sample Output:

Original array and shape:
[[[3 2 2 7 7 7 0 3]
  [5 8 4 2 9 9 3 9]
  [6 8 2 8 5 7 8 7]
  [5 2 4 0 4 9 2 5]]

 [[4 3 1 8 2 5 2 0]
  [9 1 5 8 8 5 6 5]
  [3 2 2 0 1 5 6 1]
  [5 1 9 4 2 6 9 2]]

 [[4 6 6 3 8 6 8 8]
  [3 9 2 6 3 3 1 0]
  [5 4 0 6 0 2 7 8]
  [6 3 1 8 8 1 5 7]]]
(3, 4, 8)
--------------------------------
tidex:  [0 2 2 2]
Result:
[[3 2 2 7 7 7 0 3]
 [3 9 2 6 3 3 1 0]
 [5 4 0 6 0 2 7 8]
 [6 3 1 8 8 1 5 7]]