Fix mask slicing for several classes

I should probably add more tests, but let's merge this soon if possible since otherwise the new moments stuff doesn't quite work with extracted slabs.