Resolved: add torch.testing.assert_close()

[mruberry]

Today PyTorch's test suite principally uses assertEqual(), a TestCase method, to compare tensors (and containers of tensors, scalars, and NumPy arrays).

pytorch/torch/testing/_internal/common_utils.py

Line 1229 in ea4af15

def assertEqual(self, x, y, msg: Optional[str] = None, *,

Unfortunately, this method is imperfect:

• it's a method on PyTorch's TestCase, which can make using it to compare tensors in libraries built on PyTorch tricky, and in practice there's significant use of torch.equal() to compare tensors in library tests; torch.equal() only allows comparing two tensors for bitwise equality, however
• the exact_device kwarg, which requires devices match, is still set to False for backwards compatibility; allowing tensors on different devices to compare as equal has been the source of bugs in the past
• it's called "assertEqual," but by default it actually asserts "closeness" and not "equality"

To provide a tensor comparison function suitable for both PyTorch's test suite and the test suites of libraries built on PyTorch, this RFC proposes creating a new function, torch.testing.assert_close(), with the following signature:

# Proposed testing function signature
torch.testing.assert_close(actual, expected, *, msg=None, rtol=None, atol=None, equal_nan=True, check_device=True, check_dtype=True)

# Current assertEqual signature for comparison
torch.testing.TestCase.assertEqual(self, x, y, msg=None, *, rtol, atol, equal_nan=True, exact_device=False, exact_dtype=True)

The new signature and the current signature are very similar. Their differences are:

• the new function is named assert_close to more accurately reflect its behavior and not to conflict with NumPy's testing.assert_equal name
• msg is a kwarg-only argument
• "exact_device" and "exact_dtype" have been changed to "check_device" and "check_dtype" for consistency
• the default value for "exact_device" is changed to True (note this is a TODO for assertEqual())

The behavior of this new function will be almost identical to assertEquals' behavior, except for the different default values and possibly new features like improved support for error handling.

I'd like to further propose that if and when we're happy with testing.assert_close we replace uses of assertEqual() with it in our test suite. This will create a period of a little confusion where some tests use the newer assert_close and some tests use the older assertEqual(), but our test suite has gone through more significant revisions without issue, and even assertEqual's behavior has confusingly changed over time. In fact, there are still lingering remnants of those changes left in the codebase, like assertEqualIgnoreType().

cc @ngimel @ezyang @cpuhrsch @nairbv @anjali411 @mthrok @pmeier

[mthrok]

From user perspective, the API looks good. What is the expected behavior when comparing array-like objects? such as torch.tensor([0]) == [0]?

[mruberry]

From user perspective, the API looks good. What is the expected behavior when comparing array-like objects? such as torch.tensor([0]) == [0]?

Great question! I propose we be as strict as possible and continue to declare objects of different types not equal or, more precisely, not close. So a list is never "close" to a tensor, nor would a scalar be "close" to a 0D tensor.

Allowing these cross-type comparisons can be convenient, but as a reviewer every time this feature would be used I'd have to try to figure out what the actual type was, and readability of the test suite seems important. Supporting "tensor like" objects is also tricky when working on multiple devices. For example, is zero close to torch.tensor((0,) device='cuda')? Or is it only close if check_device=False?

[ngimel]

TestCase.assertEqual allows comparison of 0d tensors with numbers.

[mruberry]

TestCase.assertEqual allows comparison of 0d tensors with numbers.

That's true, but it also defaults to exact_device=False. I'd still like to be a little stricter, and I think this is a good opportunity to make BC-breaking changes which have been challenging with assertEqual() because it's so widely used.

[NicolasHug]

Thanks for continuing this work!

I have a few questions:

• is there a plan for an additional testing.assert_equal function? I find it useful to compare int tensors which are quite frequent in torchvision. An "assert_close" would not reflect the semantic we want for integers.
• what is the status of the new utilities in testing._assert (e.g. assert_tensors_close): will they become user-facing too at some point?

[pmeier]

@NicolasHug

is there a plan for an additional testing.assert_equal function? I find it useful to compare int tensors which are quite frequent in torchvision. An "assert_close" would not reflect the semantic we want for integers.

It is already there (or will be when #55385 is landed). Plus, the assert_close function selects the default rtol / atol based in the input dtype. For all integers this is rtol = atol = 0.0.

For now the feedback we got seems to be that users are fine with that and wouldn't have much use for assert_equal. This is why we didn't promote it yet.

what is the status of the new utilities in testing._assert (e.g. assert_tensors_close): will they become user-facing too at some point?

Yes, the current plan is to get them in the upcoming 1.9 release in a prototype state. That means, they will probably not be included in the PyPI binaries, but should be accessible if you build from source or use the nightlies.

[NicolasHug]

For now the feedback we got seems to be that users are fine with that and wouldn't have much use for assert_equal

I would definitely do ;). assert_equal is more explicit than assert_close when you know the tensors you're comparing are integers. It helps "documenting" the test in a way.

Yes, the current plan is to get them in the upcoming 1.9 release in a prototype state

It seems that both assert_tensors_close and assert_close are for tensor comparison. Which one should users be using then?

[mruberry]

Hey @NicolasHug!

• is there a plan for an additional testing.assert_equal function? I find it useful to compare int tensors which are quite frequent in torchvision. An "assert_close" would not reflect the semantic we want for integers.

Great question! I agree with you it can be a convenient sugar and improve readability to use a function like testing.assert_equal() vs something like testing.assert_close(). However the worst-case workaround would be calling testing.assert_close(..., atol=0, rtol=0) or defining testing.assert_equal = partial(testing.assert_close(atol=0, rtol=0).

But yes, I think we should consider creating a public testing.assert_equal(), too!

• what is the status of the new utilities in testing._assert (e.g. assert_tensors_close): will they become user-facing too at some point?

Those are just helper functions we have no plans to publicize.

It seems that both assert_tensors_close and assert_close are for tensor comparison. Which one should users be using then?

There are some internal asserts, but just like our current compareTensors function they're just internal implementation details that we should try and keep hidden. This proposal is to just expose a testing.assert_close() function (as described above).

[antocuni]

But yes, I think we should consider creating a public testing.assert_equal(), too!

what about providing an assert_equal which works only for integer types? I.e. if you try to call it with floats you get an error which suggests to use assert_close instead. If you really, really want to compare floats by strict equality, assert_close(rtol=0, atol=0) would probably be good enough, IMHO.

Also, I don't know if I am still on time to propose changes to the name or if it's too late, but I think we could also consider assert_equal_approx or assert_approx_equal instead of assert_close:

• PRO: it has a name which is closer to the intended semantics: most of the time we are comparing something which should be mathematically equal to something else, if it weren't for the approximation errors introduced by the computers
• PRO: it is more natural for people who are used to pytest.approx. We could even think to write a plugin to allow using pytest.approx with pytorch tensors (and which would call testing.assert_close/assert_equal_approx under the hood)
• CON: it is longer to type

[nairbv]

Also, I don't know if I am still on time to propose changes to the name or if it's too late, but I think we could also consider assert_equal_approx or assert_approx_equal instead of assert_close:

Another consideration is consistency with the numpy API. Numpy has a numpy.testing.assert_almost_equal (I think we have or had a similar function) but the semantics are different (decimal places of precision). We'd want to avoid confusion between these assertions. I think the assert_close name is coming from the relationship to torch.isclose/numpy.isclose.

[mruberry]

Also, I don't know if I am still on time to propose changes to the name or if it's too late, but I think we could also consider assert_equal_approx or assert_approx_equal instead of assert_close:

Plenty of time to propose changes.

Another consideration is consistency with the numpy API. Numpy has a numpy.testing.assert_almost_equal (I think we have or had a similar function) but the semantics are different (decimal places of precision). We'd want to avoid confusion between these assertions. I think the assert_close name is coming from the relationship to torch.isclose/numpy.isclose.

Yep; I agree we don't want to use a numpy.testing name unless we implement the same function.