Regarding multi-dim indexing of numpy arrays

repl

>>> a = np.array([[1,2,3],[4,5,6]])
>>> a
array([[1, 2, 3],
       [4, 5, 6]])
>>> b = [1,2]
>>> a[:, b]
array([[2, 3],
       [5, 6]])
>>> a[np.arange(len(a)), b]
array([2, 6])