How to combine non-contiguous numpy slices

I have a numpy array with shape (M, N, N) – it’s effectively a bunch (M) of (N,N) covariance matrices. I want to be able to extract submatrices out of this with shape (M, P, P). But I’m trying to access non-contiguous indices. Here’s an example:

``````import numpy as np

# Display all the columns
np.set_printoptions(threshold=False, edgeitems=50, linewidth=200)

# Create a 6 x 6 matrix.
x = np.arange(36).reshape(6,6)

# Now make multiple copies to practice with.
y = np.array([x, x])

print(f"{y.shape=}n")

print(f"{y=}n")

# We want to extract the submatrices containting the first 2 indices
# and the last 2 indices. There are an "unknown" number of intermediate
# indices - in this example 2. Thus I'm using negative indices to get the
# last two indicies.

s = np.array([[0, 1] + [-2, -1]])

subset = y[:, s.T, s]
print(f"{subset=}n")

# Now try it with numpy slices. This approach doesn't work
first_slice = np.s_[0:2]
second_slice = np.s_[-2:]
combined_slice = np.r_[first_slice, second_slice]

subset = y[:, combined_slice, combined_slice]
print(subset)
``````

That produces the following output:

``````y.shape=(2, 6, 6)

y=array([[[ 0,  1,  2,  3,  4,  5],
[ 6,  7,  8,  9, 10, 11],
[12, 13, 14, 15, 16, 17],
[18, 19, 20, 21, 22, 23],
[24, 25, 26, 27, 28, 29],
[30, 31, 32, 33, 34, 35]],

[[ 0,  1,  2,  3,  4,  5],
[ 6,  7,  8,  9, 10, 11],
[12, 13, 14, 15, 16, 17],
[18, 19, 20, 21, 22, 23],
[24, 25, 26, 27, 28, 29],
[30, 31, 32, 33, 34, 35]]])

subset=array([[[ 0,  1,  4,  5],
[ 6,  7, 10, 11],
[24, 25, 28, 29],
[30, 31, 34, 35]],

[[ 0,  1,  4,  5],
[ 6,  7, 10, 11],
[24, 25, 28, 29],
[30, 31, 34, 35]]])

[[0 7]
[0 7]]
``````

Using array based indexing works, but it seems like there should be a way to do this with slice objects too. Any tips would be appreciated. Thanks!