Function _assert_tensors_equal

tests/test_modeling_bart.py:534–546

If tensors not close, or a and b arent both tensors, raise a nice Assertion error.

(a, b, atol=1e-12, prefix="")

def _assert_tensors_equal(a, b, atol=1e-12, prefix=""):
    """If tensors not close, or a and b arent both tensors, raise a nice Assertion error."""
    if a is None and b is None:
        return True
    try:
        if torch.allclose(a, b, atol=atol):
            return True
        raise
    except Exception:
        msg = "{} != {}".format(a, b)
        if prefix:
            msg = prefix + ": " + msg
        raise AssertionError(msg)


def _long_tensor(tok_lst):