reference/generated/tensordict.assert_allclose_td

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

assert_allclose_td

tensordict.assert_allclose_td(actual: T, expected: T, rtol: float | None = None, atol: float | None = None, equal_nan: bool = True, intersection: bool = False, msg: str = '', prefix: NestedKey = ()) → bool

Asserts that two tensordicts have the same keys and close values.

Each pair of tensors is compared with torch.testing.assert_close(), which also checks that the dtypes, shapes and devices match. Nested tensordicts are compared entry by entry, and two lazy stacks are compared stack by stack. Non-tensor data is not compared. Tensorclasses are accepted.

assert_allclose_td is another name for this function.

Parameters:

• actual (TensorDictBase) – the tensordict to check.

• expected (TensorDictBase) – the tensordict with the expected values.

• rtol (float, optional) – the relative tolerance. Defaults to None, which uses the default of torch.testing.assert_close() for the dtype.

• atol (float, optional) – the absolute tolerance. Defaults to None, which uses the default of torch.testing.assert_close() for the dtype.

• equal_nan (bool, optional) – if True, two NaN values are equal. Defaults to True.

• intersection (bool, optional) – if True, compares only the keys that both tensordicts have. If False, different key sets raise a KeyError. Defaults to False.

• msg (str, optional) – a message appended to the error message. Defaults to "".

• prefix (NestedKey, optional) – a key prepended to the keys in the error messages. Defaults to ().

Returns:

True if the check passes.

Raises:

• TypeError – if actual or expected is not a tensordict or a tensorclass.

• KeyError – if the keys differ and intersection=False.

• AssertionError – if two values are not close. The message names the key and the root mean squared error, followed by torch.testing.assert_close()’s description of the mismatch.

• RuntimeError – if the shapes of two tensors cannot be broadcast together.

Examples

>>> import torch
>>> from tensordict import TensorDict, assert_close
>>> a = TensorDict(x=torch.zeros(3), y=torch.ones(3))
>>> assert_close(a, a.clone())
True
>>> b = TensorDict(x=torch.full((3,), 0.01), y=torch.ones(3))
>>> assert_close(a, b, atol=0.1, rtol=0)
True
>>> assert_close(a, b)  
Traceback (most recent call last):
    ...
AssertionError: key ('x',) does not match, got mse = 0.0100
Tensor-likes are not close!
...