Function assert_allclose

tests/test_utils.py:123–163

Assert that types and all values of two data objects are close. Args: actual: Pytorch Tensor or numpy array for comparison. desired: Pytorch Tensor or numpy array to compare against. type_test: whether to test that `actual` and `desired` are both numpy arrays or tor

(
    actual: NdarrayOrTensor,
    desired: NdarrayOrTensor,
    type_test: bool | str = True,
    device_test: bool = False,
    *args,
    **kwargs,
)

def assert_allclose(
    actual: NdarrayOrTensor,
    desired: NdarrayOrTensor,
    type_test: bool | str = True,
    device_test: bool = False,
    *args,
    **kwargs,
):
    """
    Assert that types and all values of two data objects are close.

    Args:
        actual: Pytorch Tensor or numpy array for comparison.
        desired: Pytorch Tensor or numpy array to compare against.
        type_test: whether to test that `actual` and `desired` are both numpy arrays or torch tensors.
            if type_test == "tensor", it checks whether the `actual` is a torch.tensor or metatensor according to
            `get_track_meta`.
        device_test: whether to test the device property.
        args: extra arguments to pass on to `np.testing.assert_allclose`.
        kwargs: extra arguments to pass on to `np.testing.assert_allclose`.


    """
    if isinstance(type_test, str) and type_test == "tensor":
        if get_track_meta():
            np.testing.assert_equal(isinstance(actual, MetaTensor), True, "must be a MetaTensor")
        else:
            np.testing.assert_equal(
                isinstance(actual, torch.Tensor) and not isinstance(actual, MetaTensor), True, "must be a torch.Tensor"
            )
    elif type_test:
        # check both actual and desired are of the same type
        np.testing.assert_equal(isinstance(actual, np.ndarray), isinstance(desired, np.ndarray), "numpy type")
        np.testing.assert_equal(isinstance(actual, torch.Tensor), isinstance(desired, torch.Tensor), "torch type")

    if isinstance(desired, torch.Tensor) or isinstance(actual, torch.Tensor):
        if device_test:
            np.testing.assert_equal(str(actual.device), str(desired.device), "torch device check")  # type: ignore
        actual = actual.detach().cpu().numpy() if isinstance(actual, torch.Tensor) else actual
        desired = desired.detach().cpu().numpy() if isinstance(desired, torch.Tensor) else desired
    np.testing.assert_allclose(actual, desired, *args, **kwargs)


@contextmanager