Apply 1d-mask on numpy 3d-array

I have the following 3d-numpy.ndarray:

import numpy as np
X = np.array([
    [[0.0, 0.4, 0.6, 0.0, 0.0],
     [0.6, 0.0, 0.0, 0.0, 0.0],
     [0.4, 0.0, 0.0, 0.0, 0.0],
     [0.0, 0.6, 0.0, 1.0, 0.0],
     [0.0, 0.0, 0.4, 0.0, 1.0]],
    [[0.1, 0.5, 0.4, 0.0, 0.0],
     [0.6, 0.0, 0.0, 0.0, 0.0],
     [0.2, 0.0, 0.0, 0.0, 0.0],
     [0.1, 0.6, 0.0, 1.0, 0.0],
     [0.0, 0.0, 0.4, 0.0, 1.0]]
])

I want a new array where all the rows and columns are dropped where the diagnonal is equal to 1.

idx = np.diag(X[0]) == 1          # for my implementation it is sufficient to look at X[0]

Important to note is that X.shape[1] == X.shape[2], so I try to use the mask as follows

Y = X[:, ~idx, ~idx]

The above returns something different than my desired output:

 [[0.0, 0.4, 0.6],
  [0.6, 0.0, 0.0],
  [0.4, 0.0, 0.0]],
 [[0.1, 0.5, 0.4],
  [0.6, 0.0, 0.0],
  [0.2, 0.0, 0.0]]

Please advice

Answer

Probably you can try

X[:, ~idx,:][:,:, ~idx]

or np.ix_

X[:, *np.ix_(~idx, ~idx)]

which gives

array([[[0. , 0.4, 0.6],
        [0.6, 0. , 0. ],
        [0.4, 0. , 0. ]],

       [[0.1, 0.5, 0.4],
        [0.6, 0. , 0. ],
        [0.2, 0. , 0. ]]])