The case of indexing with `arr[bool_matrix] -> arr[bool_matrix.nonzero()]` needs to be dealt with in this function, I think.