Add basic tests for drop_dims