Torch

recursive_assert_approx_equal(x, y, atol=0.0001, rtol=0.0001)

Assert that all tensors in a nested structure are approximately equal.

Source code in bionemo/testing/torch.py

def recursive_assert_approx_equal(x, y, atol=1e-4, rtol=1e-4):
    """Assert that all tensors in a nested structure are approximately equal."""
    if isinstance(x, torch.Tensor):
        torch.testing.assert_close(x, y, atol=atol, rtol=rtol)
    elif isinstance(x, np.ndarray):
        np.testing.assert_allclose(x, y, atol=atol, rtol=rtol)
    elif isinstance(x, (list, tuple)):
        assert len(x) == len(y), f"Length mismatch: {len(x)} vs {len(y)}"
        for x_item, y_item in zip(x, y):
            recursive_assert_approx_equal(x_item, y_item, atol=atol, rtol=rtol)
    elif isinstance(x, dict):
        assert x.keys() == y.keys()
        for key in x:
            recursive_assert_approx_equal(x[key], y[key], atol=atol, rtol=rtol)
    else:
        assert x == y

recursive_detach(x)

Detach all tensors in a nested structure.

Source code in bionemo/testing/torch.py

def recursive_detach(x):
    """Detach all tensors in a nested structure."""
    if isinstance(x, torch.Tensor):
        return x.detach().cpu()
    elif isinstance(x, (list, tuple)):
        return type(x)(recursive_detach(item) for item in x)
    elif isinstance(x, dict):
        return {key: recursive_detach(value) for key, value in x.items()}
    else:
        return x