I'm not sure if it's necessary, but I think you should use np.testing.assert_equals here. @ahojnnes?