Masking along one dimension

This topic has been discussed but wasn’t finished:

I ran into a similar problem. In my case I have a tensor with [60000, 100, 400], that I need to do some operation along the third dimension. I don’t need to do any operation on most of the entries along third dimension. I was able to produce a mask of the same size that I hope reduce the tensor to something like [60000, 100, 150], but I was only able to either flatten the entire array, or keep the original dimension and make masked values zero, either way doesn’t help.

The number of unmasked values for the third dimension also varies from a few to roughly 150… for this I’m thinking making the extra values NaN would be fine.