选择数组中向量的分量

Question

I have a tensor with shape (7, 2, 3) I want to select one of the two row vectors from each of the 7 2x3 matrices, ie

[
 [[0, 0, 0],
  [1, 1, 1]],

 [[0, 0, 0],
  [1, 1, 1]],
 ...x7
]

to

a = [
 [0, 0, 0],
 [0, 0, 0],
 [0, 0, 0]
 ...x7
]
b = [
 [1, 1, 1],
 [1, 1, 1],
 [1, 1, 1]
 ...x7
]

each with shape (7, 3).

How can I do this without reshape ? (I find reshape to be kind of confusing when some dimensions are the same).

I also know of

np.array(map(lambda item: item[0], x)))

but I would like a more concise way if there is one.

Answer 1

just use looped indexing: data[:, i, :] where i loops from 0 through 1

import numpy as np

a = np.array([
 [[0, 0, 0],
  [1, 1, 1]],

 [[0, 0, 0],
  [1, 1, 1]]
])
print(a[:, 1, :])

will produce

[[1 1 1]
 [1 1 1]]