I usually prefer `np.testing.assert_array_equal` to `assert np.all` since it provides better error messages -- is there a reason why some of the tests use np.all here? If not, can you switch them?