jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_3_channel_tensor_to_pil_image(self, expected_mode):
        img_data = torch.Tensor(3, 4, 4).uniform_()
        expected_output = img_data.mul(255).int().float().div(255)

        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "RGB"  # default should assume RGB
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode
        split = img.split()
        for i in range(3):
            torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
//  [python]
def test_4_channel_tensor_to_pil_image(self, expected_mode):
        img_data = torch.Tensor(4, 4, 4).uniform_()
        expected_output = img_data.mul(255).int().float().div(255)

        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "RGBA"  # default should assume RGBA
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode

        split = img.split()
        for i in range(4):
            torch.testing.assert_allclose(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
//  [python]
def test_4_channel_tensor_to_pil_image(self, expected_mode):
        img_data = torch.Tensor(4, 4, 4).uniform_()
        expected_output = img_data.mul(255).int().float().div(255)

        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "RGBA"  # default should assume RGBA
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode

        split = img.split()
        for i in range(4):
            torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy())
//  [python]
def test_4_channel_ndarray_to_pil_image(self, expected_mode):
        img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy()

        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "RGBA"  # default should assume RGBA
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode
        split = img.split()
        for i in range(4):
            torch.testing.assert_allclose(img_data[:, :, i], np.asarray(split[i]))
//  [python]
def test_4_channel_ndarray_to_pil_image(self, expected_mode):
        img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy()

        if expected_mode is None:
            img = transforms.ToPILImage()(img_data)
            assert img.mode == "RGBA"  # default should assume RGBA
        else:
            img = transforms.ToPILImage(mode=expected_mode)(img_data)
            assert img.mode == expected_mode
        split = img.split()
        for i in range(4):
            torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]))
//  [python]
def test_adjust_brightness():
    x_shape = [2, 2, 3]
    x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")

    # test 0
    y_pil = F.adjust_brightness(x_pil, 1)
    y_np = np.array(y_pil)
    t.assert_close(y_np, x_np)

    # test 1
    y_pil = F.adjust_brightness(x_pil, 0.5)
    y_np = np.array(y_pil)
    y_ans = [0, 2, 6, 27, 67, 113, 18, 4, 117, 45, 127, 0]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_close(y_np, y_ans)

    # test 2
    y_pil = F.adjust_brightness(x_pil, 2)
    y_np = np.array(y_pil)
    y_ans = [0, 10, 26, 108, 255, 255, 74, 16, 255, 180, 255, 2]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_close(y_np, y_ans)
//  [python]
def test_adjust_brightness():
    x_shape = [2, 2, 3]
    x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")

    # test 0
    y_pil = F.adjust_brightness(x_pil, 1)
    y_np = np.array(y_pil)
    torch.testing.assert_close(y_np, x_np)

    # test 1
    y_pil = F.adjust_brightness(x_pil, 0.5)
    y_np = np.array(y_pil)
    y_ans = [0, 2, 6, 27, 67, 113, 18, 4, 117, 45, 127, 0]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)

    # test 2
    y_pil = F.adjust_brightness(x_pil, 2)
    y_np = np.array(y_pil)
    y_ans = [0, 10, 26, 108, 255, 255, 74, 16, 255, 180, 255, 2]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)
//  [python]
def test_adjust_contrast():
    x_shape = [2, 2, 3]
    x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")

    # test 0
    y_pil = F.adjust_contrast(x_pil, 1)
    y_np = np.array(y_pil)
    t.assert_close(y_np, x_np)

    # test 1
    y_pil = F.adjust_contrast(x_pil, 0.5)
    y_np = np.array(y_pil)
    y_ans = [43, 45, 49, 70, 110, 156, 61, 47, 160, 88, 170, 43]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_close(y_np, y_ans)

    # test 2
    y_pil = F.adjust_contrast(x_pil, 2)
    y_np = np.array(y_pil)
    y_ans = [0, 0, 0, 22, 184, 255, 0, 0, 255, 94, 255, 0]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_close(y_np, y_ans)
//  [python]
def test_adjust_contrast():
    x_shape = [2, 2, 3]
    x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")

    # test 0
    y_pil = F.adjust_contrast(x_pil, 1)
    y_np = np.array(y_pil)
    torch.testing.assert_close(y_np, x_np)

    # test 1
    y_pil = F.adjust_contrast(x_pil, 0.5)
    y_np = np.array(y_pil)
    y_ans = [43, 45, 49, 70, 110, 156, 61, 47, 160, 88, 170, 43]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)

    # test 2
    y_pil = F.adjust_contrast(x_pil, 2)
    y_np = np.array(y_pil)
    y_ans = [0, 0, 0, 22, 184, 255, 0, 0, 255, 94, 255, 0]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)
//  [python]
def test_normalize_3d_tensor():
    torch.manual_seed(28)
    n_channels = 3
    img_size = 10
    mean = torch.rand(n_channels)
    std = torch.rand(n_channels)
    img = torch.rand(n_channels, img_size, img_size)
    target = F.normalize(img, mean, std)

    mean_unsqueezed = mean.view(-1, 1, 1)
    std_unsqueezed = std.view(-1, 1, 1)
    result1 = F.normalize(img, mean_unsqueezed, std_unsqueezed)
    result2 = F.normalize(
        img, mean_unsqueezed.repeat(1, img_size, img_size), std_unsqueezed.repeat(1, img_size, img_size)
    )
    torch.testing.assert_allclose(target, result1)
    torch.testing.assert_allclose(target, result2)
//  [python]
def test_normalize_3d_tensor():
    torch.manual_seed(28)
    n_channels = 3
    img_size = 10
    mean = torch.rand(n_channels)
    std = torch.rand(n_channels)
    img = torch.rand(n_channels, img_size, img_size)
    target = F.normalize(img, mean, std)

    mean_unsqueezed = mean.view(-1, 1, 1)
    std_unsqueezed = std.view(-1, 1, 1)
    result1 = F.normalize(img, mean_unsqueezed, std_unsqueezed)
    result2 = F.normalize(
        img, mean_unsqueezed.repeat(1, img_size, img_size), std_unsqueezed.repeat(1, img_size, img_size)
    )
    torch.testing.assert_close(target, result1)
    torch.testing.assert_close(target, result2)
//  [python]
def test_adjust_hue():
    x_shape = [2, 2, 3]
    x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")

    with pytest.raises(ValueError):
        F.adjust_hue(x_pil, -0.7)
        F.adjust_hue(x_pil, 1)

    # test 0: almost same as x_data but not exact.
    # probably because hsv <-> rgb floating point ops
    y_pil = F.adjust_hue(x_pil, 0)
    y_np = np.array(y_pil)
    y_ans = [0, 5, 13, 54, 139, 226, 35, 8, 234, 91, 255, 1]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_allclose(y_np, y_ans)

    # test 1
    y_pil = F.adjust_hue(x_pil, 0.25)
    y_np = np.array(y_pil)
    y_ans = [13, 0, 12, 224, 54, 226, 234, 8, 99, 1, 222, 255]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_allclose(y_np, y_ans)

    # test 2
    y_pil = F.adjust_hue(x_pil, -0.25)
    y_np = np.array(y_pil)
    y_ans = [0, 13, 2, 54, 226, 58, 8, 234, 152, 255, 43, 1]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_allclose(y_np, y_ans)
//  [python]
def test_adjust_hue():
    x_shape = [2, 2, 3]
    x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")

    with pytest.raises(ValueError):
        F.adjust_hue(x_pil, -0.7)
        F.adjust_hue(x_pil, 1)

    # test 0: almost same as x_data but not exact.
    # probably because hsv <-> rgb floating point ops
    y_pil = F.adjust_hue(x_pil, 0)
    y_np = np.array(y_pil)
    y_ans = [0, 5, 13, 54, 139, 226, 35, 8, 234, 91, 255, 1]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)

    # test 1
    y_pil = F.adjust_hue(x_pil, 0.25)
    y_np = np.array(y_pil)
    y_ans = [13, 0, 12, 224, 54, 226, 234, 8, 99, 1, 222, 255]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)

    # test 2
    y_pil = F.adjust_hue(x_pil, -0.25)
    y_np = np.array(y_pil)
    y_ans = [0, 13, 2, 54, 226, 58, 8, 234, 152, 255, 43, 1]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)
//  [python]
def test_adjust_gamma():
    x_shape = [2, 2, 3]
    x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")

    # test 0
    y_pil = F.adjust_gamma(x_pil, 1)
    y_np = np.array(y_pil)
    t.assert_close(y_np, x_np)

    # test 1
    y_pil = F.adjust_gamma(x_pil, 0.5)
    y_np = np.array(y_pil)
    y_ans = [0, 35, 57, 117, 186, 241, 97, 45, 245, 152, 255, 16]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_close(y_np, y_ans)

    # test 2
    y_pil = F.adjust_gamma(x_pil, 2)
    y_np = np.array(y_pil)
    y_ans = [0, 0, 0, 11, 71, 201, 5, 0, 215, 31, 255, 0]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_close(y_np, y_ans)
//  [python]
def test_adjust_gamma():
    x_shape = [2, 2, 3]
    x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")

    # test 0
    y_pil = F.adjust_gamma(x_pil, 1)
    y_np = np.array(y_pil)
    torch.testing.assert_close(y_np, x_np)

    # test 1
    y_pil = F.adjust_gamma(x_pil, 0.5)
    y_np = np.array(y_pil)
    y_ans = [0, 35, 57, 117, 186, 241, 97, 45, 245, 152, 255, 16]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)

    # test 2
    y_pil = F.adjust_gamma(x_pil, 2)
    y_np = np.array(y_pil)
    y_ans = [0, 0, 0, 11, 71, 201, 5, 0, 215, 31, 255, 0]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)
//  [python]
def test_random_grayscale_with_grayscale_input():
    transform = transforms.RandomGrayscale(p=1.0)

    image_tensor = torch.randint(0, 256, (1, 16, 16), dtype=torch.uint8)
    output_tensor = transform(image_tensor)
    torch.testing.assert_allclose(output_tensor, image_tensor)

    image_pil = F.to_pil_image(image_tensor)
    output_pil = transform(image_pil)
    torch.testing.assert_allclose(F.pil_to_tensor(output_pil), image_tensor)
//  [python]
def test_random_grayscale_with_grayscale_input():
    transform = transforms.RandomGrayscale(p=1.0)

    image_tensor = torch.randint(0, 256, (1, 16, 16), dtype=torch.uint8)
    output_tensor = transform(image_tensor)
    torch.testing.assert_close(output_tensor, image_tensor)

    image_pil = F.to_pil_image(image_tensor)
    output_pil = transform(image_pil)
    torch.testing.assert_close(F.pil_to_tensor(output_pil), image_tensor)
//  [python]
def test_random_crop():
    height = random.randint(10, 32) * 2
    width = random.randint(10, 32) * 2
    oheight = random.randint(5, (height - 2) // 2) * 2
    owidth = random.randint(5, (width - 2) // 2) * 2
    img = torch.ones(3, height, width, dtype=torch.uint8)
    result = transforms.Compose(
        [
            transforms.ToPILImage(),
            transforms.RandomCrop((oheight, owidth)),
            transforms.PILToTensor(),
        ]
    )(img)
    assert result.size(1) == oheight
    assert result.size(2) == owidth

    padding = random.randint(1, 20)
    result = transforms.Compose(
        [
            transforms.ToPILImage(),
            transforms.RandomCrop((oheight, owidth), padding=padding),
            transforms.PILToTensor(),
        ]
    )(img)
    assert result.size(1) == oheight
    assert result.size(2) == owidth

    result = transforms.Compose(
        [transforms.ToPILImage(), transforms.RandomCrop((height, width)), transforms.PILToTensor()]
    )(img)
    assert result.size(1) == height
    assert result.size(2) == width
    torch.allclose(result, img)

    result = transforms.Compose(
        [
            transforms.ToPILImage(),
            transforms.RandomCrop((height + 1, width + 1), pad_if_needed=True),
            transforms.PILToTensor(),
        ]
    )(img)
    assert result.size(1) == height + 1
    assert result.size(2) == width + 1

    t = transforms.RandomCrop(33)
    img = torch.ones(3, 32, 32)
    with pytest.raises(ValueError, match=r"Required crop size .+ is larger than input image size .+"):
        t(img)
//  [python]
def test_random_crop():
    height = random.randint(10, 32) * 2
    width = random.randint(10, 32) * 2
    oheight = random.randint(5, (height - 2) // 2) * 2
    owidth = random.randint(5, (width - 2) // 2) * 2
    img = torch.ones(3, height, width, dtype=torch.uint8)
    result = transforms.Compose(
        [
            transforms.ToPILImage(),
            transforms.RandomCrop((oheight, owidth)),
            transforms.PILToTensor(),
        ]
    )(img)
    assert result.size(1) == oheight
    assert result.size(2) == owidth

    padding = random.randint(1, 20)
    result = transforms.Compose(
        [
            transforms.ToPILImage(),
            transforms.RandomCrop((oheight, owidth), padding=padding),
            transforms.PILToTensor(),
        ]
    )(img)
    assert result.size(1) == oheight
    assert result.size(2) == owidth

    result = transforms.Compose(
        [transforms.ToPILImage(), transforms.RandomCrop((height, width)), transforms.PILToTensor()]
    )(img)
    assert result.size(1) == height
    assert result.size(2) == width
    torch.testing.assert_close(result, img)

    result = transforms.Compose(
        [
            transforms.ToPILImage(),
            transforms.RandomCrop((height + 1, width + 1), pad_if_needed=True),
            transforms.PILToTensor(),
        ]
    )(img)
    assert result.size(1) == height + 1
    assert result.size(2) == width + 1

    t = transforms.RandomCrop(33)
    img = torch.ones(3, 32, 32)
    with pytest.raises(ValueError, match=r"Required crop size .+ is larger than input image size .+"):
        t(img)
//  [python]
def test_linear_transformation():
    num_samples = 1000
    x = torch.randn(num_samples, 3, 10, 10)
    flat_x = x.view(x.size(0), x.size(1) * x.size(2) * x.size(3))
    # compute principal components
    sigma = torch.mm(flat_x.t(), flat_x) / flat_x.size(0)
    u, s, _ = np.linalg.svd(sigma.numpy())
    zca_epsilon = 1e-10  # avoid division by 0
    d = torch.Tensor(np.diag(1.0 / np.sqrt(s + zca_epsilon)))
    u = torch.Tensor(u)
    principal_components = torch.mm(torch.mm(u, d), u.t())
    mean_vector = torch.sum(flat_x, dim=0) / flat_x.size(0)
    # initialize whitening matrix
    whitening = transforms.LinearTransformation(principal_components, mean_vector)
    # estimate covariance and mean using weak law of large number
    num_features = flat_x.size(1)
    cov = 0.0
    mean = 0.0
    for i in x:
        xwhite = whitening(i)
        xwhite = xwhite.view(1, -1).numpy()
        cov += np.dot(xwhite, xwhite.T) / num_features
        mean += np.sum(xwhite) / num_features
    # if rtol for std = 1e-3 then rtol for cov = 2e-3 as std**2 = cov
    t.assert_close(
        cov / num_samples, np.identity(1), rtol=2e-3, atol=1e-8, check_dtype=False, msg="cov not close to 1"
    )
    t.assert_close(
        mean / num_samples, 0, rtol=1e-3, atol=1e-8, check_dtype=False, msg="mean not close to 0"
    )

    # Checking if LinearTransformation can be printed as string
    whitening.__repr__()
//  [python]
def test_linear_transformation():
    num_samples = 1000
    x = torch.randn(num_samples, 3, 10, 10)
    flat_x = x.view(x.size(0), x.size(1) * x.size(2) * x.size(3))
    # compute principal components
    sigma = torch.mm(flat_x.t(), flat_x) / flat_x.size(0)
    u, s, _ = np.linalg.svd(sigma.numpy())
    zca_epsilon = 1e-10  # avoid division by 0
    d = torch.Tensor(np.diag(1.0 / np.sqrt(s + zca_epsilon)))
    u = torch.Tensor(u)
    principal_components = torch.mm(torch.mm(u, d), u.t())
    mean_vector = torch.sum(flat_x, dim=0) / flat_x.size(0)
    # initialize whitening matrix
    whitening = transforms.LinearTransformation(principal_components, mean_vector)
    # estimate covariance and mean using weak law of large number
    num_features = flat_x.size(1)
    cov = 0.0
    mean = 0.0
    for i in x:
        xwhite = whitening(i)
        xwhite = xwhite.view(1, -1).numpy()
        cov += np.dot(xwhite, xwhite.T) / num_features
        mean += np.sum(xwhite) / num_features
    # if rtol for std = 1e-3 then rtol for cov = 2e-3 as std**2 = cov
    torch.testing.assert_close(
        cov / num_samples, np.identity(1), rtol=2e-3, atol=1e-8, check_dtype=False, msg="cov not close to 1"
    )
    torch.testing.assert_close(
        mean / num_samples, 0, rtol=1e-3, atol=1e-8, check_dtype=False, msg="mean not close to 0"
    )

    # Checking if LinearTransformation can be printed as string
    whitening.__repr__()
//  [python]
def _test_against_reference(self, inputs, reference, eps, dtype):
        q, k, v, rpb, kernel_size, dilation, is_causal = inputs
        kernel_size, dilation, is_causal = check_args(kernel_size, dilation, is_causal)
        out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            out = na2d(
                q_,
                k_,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
                rpb=rpb_,
            )

            torch.testing.assert_allclose(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def _test_against_reference(self, inputs, reference, eps, dtype):
        q, k, v, rpb, kernel_size, dilation, is_causal = inputs
        kernel_size, dilation, is_causal = check_args(kernel_size, dilation, is_causal)
        out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            out = na2d(
                q_,
                k_,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
                rpb=rpb_,
            )

            torch.testing.assert_close(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def _test_against_sdpa(self, inputs, reference, eps, dtype):
        q, k, v = inputs
        kernel_size = (q.shape[1], q.shape[2])
        kernel_size, dilation, is_causal = check_args(kernel_size, 1, False)
        out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            out = na2d(
                q_,
                k_,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
            )

            torch.testing.assert_allclose(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def _test_against_sdpa(self, inputs, reference, eps, dtype):
        q, k, v = inputs
        kernel_size = (q.shape[1], q.shape[2])
        kernel_size, dilation, is_causal = check_args(kernel_size, 1, False)
        out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            out = na2d(
                q_,
                k_,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
            )

            torch.testing.assert_close(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def test_torch_mix_cuda_cpu():
    model = CopyTensorModule()
    x = torch.randn(3, 4, device='cpu')
    y = model(x)

    model_opt = torch.compile(model, backend='hidet')
    y1 = model_opt(x)

    torch.testing.assert_allclose(y, y1, rtol=0.0, atol=0.0)
//  [python]
def test_torch_mix_cuda_cpu():
    model = CopyTensorModule()
    x = torch.randn(3, 4, device='cpu')
    y = model(x)

    model_opt = torch.compile(model, backend='hidet')
    y1 = model_opt(x)

    torch.testing.assert_close(y, y1, rtol=0.0, atol=0.0)
//  [python]
def test_yolo_transform():
    transform = YOLOTransform(300, 500)
    images = [torch.rand(3, 200, 300), torch.rand(3, 200, 200)]
    annotations = [
        {"boxes": torch.randint(180, (3, 4)), "labels": torch.randint(80, (3,))},
        {"boxes": torch.randint(180, (2, 4)), "labels": torch.randint(80, (2,))},
    ]
    annotations_copy = copy.deepcopy(annotations)
    samples, targets = transform(images, annotations)
    assert isinstance(samples, NestedTensor)
    assert targets.shape[1] == 6

    # Test annotations after transformation
    t.assert_close(annotations[0]["boxes"], annotations_copy[0]["boxes"], rtol=0, atol=0)
    t.assert_close(annotations[1]["boxes"], annotations_copy[1]["boxes"], rtol=0, atol=0)
//  [python]
def test_yolo_transform():
    transform = YOLOTransform(300, 500)
    images = [torch.rand(3, 200, 300), torch.rand(3, 200, 200)]
    annotations = [
        {"boxes": torch.randint(180, (3, 4)), "labels": torch.randint(80, (3,))},
        {"boxes": torch.randint(180, (2, 4)), "labels": torch.randint(80, (2,))},
    ]
    annotations_copy = copy.deepcopy(annotations)
    samples, targets = transform(images, annotations)
    assert isinstance(samples, NestedTensor)
    assert targets.shape[1] == 6

    # Test annotations after transformation
    torch.testing.assert_close(annotations[0]["boxes"], annotations_copy[0]["boxes"], rtol=0, atol=0)
    torch.testing.assert_close(annotations[1]["boxes"], annotations_copy[1]["boxes"], rtol=0, atol=0)
//  [python]
def test_default_2d_filter():

    def reference_filter():
        filt = torch.FloatTensor([1, 2, 1])
        filt = torch.outer(filt, filt)
        filt *= 1. / filt.sum()
        filt = torch.Tensor(filt)
        return filt.view(1, 1, *filt.shape)

    torch.testing.assert_allclose(
        blurpool.blurpool_layers._default_2d_filter(),  # type: ignore
        reference_filter(),
    )
//  [python]
def test_default_2d_filter():

    def reference_filter():
        filt = torch.FloatTensor([1, 2, 1])
        filt = torch.outer(filt, filt)
        filt *= 1. / filt.sum()
        filt = torch.Tensor(filt)
        return filt.view(1, 1, *filt.shape)

    torch.testing.assert_close(
        blurpool.blurpool_layers._default_2d_filter(),  # type: ignore
        reference_filter(),
    )
//  [python]
def _test_large_cum_fn_helper(self, x, fn):
        expected = fn(x.cpu().float())
        actual = fn(x).cpu().float()
        # Avoid self.assertEqual to save memory.
        self.assertEqual(expected, actual)
//  [python]
def _test_large_cum_fn_helper(self, x, fn):
        expected = fn(x.cpu().float())
        actual = fn(x).cpu().float()
        # Avoid self.assertEqual to save memory.
        torch.testing.assert_close(expected, actual)
//  [python]
def test_blurpool_blurconv2d_params_match_original_params():
    conv2d = torch.nn.Conv2d(16, 32, 3, stride=1, bias=True)
    blurconv = blurpool.BlurConv2d.from_conv2d(conv2d)
    torch.testing.assert_allclose(blurconv.conv.weight, conv2d.weight)
    torch.testing.assert_allclose(blurconv.conv.bias, conv2d.bias)
    assert blurconv.conv.weight.requires_grad
    assert blurconv.conv.bias is not None
    assert blurconv.conv.bias.requires_grad
//  [python]
def test_blurpool_blurconv2d_params_match_original_params():
    conv2d = torch.nn.Conv2d(16, 32, 3, stride=1, bias=True)
    blurconv = blurpool.BlurConv2d.from_conv2d(conv2d)
    torch.testing.assert_close(blurconv.conv.weight, conv2d.weight)
    torch.testing.assert_close(blurconv.conv.bias, conv2d.bias)
    assert blurconv.conv.weight.requires_grad
    assert blurconv.conv.bias is not None
    assert blurconv.conv.bias.requires_grad
//  [python]
def test_sobolengine_continuing(self, scramble: bool = False):
        ref_sample = self._sobol_reference_samples(scramble=scramble)
        engine = torch.quasirandom.SobolEngine(2, scramble=scramble, seed=123456)
        n_half = len(ref_sample) // 2
        _ = engine.draw(n=n_half)
        sample = engine.draw(n=n_half)
        self.assertEqual(sample, ref_sample[n_half:])
//  [python]
def test_sobolengine_continuing(self, scramble: bool = False):
        ref_sample = self._sobol_reference_samples(scramble=scramble)
        engine = torch.quasirandom.SobolEngine(2, scramble=scramble, seed=123456)
        n_half = len(ref_sample) // 2
        _ = engine.draw(n=n_half)
        sample = engine.draw(n=n_half)
        torch.testing.assert_close(sample, ref_sample[n_half:])
//  [python]
def test_adjust_sharpness():
    x_shape = [4, 4, 3]
    x_data = [
        75,
        121,
        114,
        105,
        97,
        107,
        105,
        32,
        66,
        111,
        117,
        114,
        99,
        104,
        97,
        0,
        0,
        65,
        108,
        101,
        120,
        97,
        110,
        100,
        101,
        114,
        32,
        86,
        114,
        121,
        110,
        105,
        111,
        116,
        105,
        115,
        0,
        0,
        73,
        32,
        108,
        111,
        118,
        101,
        32,
        121,
        111,
        117,
    ]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")

    # test 0
    y_pil = F.adjust_sharpness(x_pil, 1)
    y_np = np.array(y_pil)
    t.assert_close(y_np, x_np)

    # test 1
    y_pil = F.adjust_sharpness(x_pil, 0.5)
    y_np = np.array(y_pil)
    y_ans = [
        75,
        121,
        114,
        105,
        97,
        107,
        105,
        32,
        66,
        111,
        117,
        114,
        99,
        104,
        97,
        30,
        30,
        74,
        103,
        96,
        114,
        97,
        110,
        100,
        101,
        114,
        32,
        81,
        103,
        108,
        102,
        101,
        107,
        116,
        105,
        115,
        0,
        0,
        73,
        32,
        108,
        111,
        118,
        101,
        32,
        121,
        111,
        117,
    ]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_close(y_np, y_ans)

    # test 2
    y_pil = F.adjust_sharpness(x_pil, 2)
    y_np = np.array(y_pil)
    y_ans = [
        75,
        121,
        114,
        105,
        97,
        107,
        105,
        32,
        66,
        111,
        117,
        114,
        99,
        104,
        97,
        0,
        0,
        46,
        118,
        111,
        132,
        97,
        110,
        100,
        101,
        114,
        32,
        95,
        135,
        146,
        126,
        112,
        119,
        116,
        105,
        115,
        0,
        0,
        73,
        32,
        108,
        111,
        118,
        101,
        32,
        121,
        111,
        117,
    ]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    t.assert_close(y_np, y_ans)

    # test 3
    x_shape = [2, 2, 3]
    x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")
    x_th = torch.tensor(x_np.transpose(2, 0, 1))
    y_pil = F.adjust_sharpness(x_pil, 2)
    y_np = np.array(y_pil).transpose(2, 0, 1)
    y_th = F.adjust_sharpness(x_th, 2)
    t.assert_close(y_np, y_th.numpy())
//  [python]
def test_adjust_sharpness():
    x_shape = [4, 4, 3]
    x_data = [
        75,
        121,
        114,
        105,
        97,
        107,
        105,
        32,
        66,
        111,
        117,
        114,
        99,
        104,
        97,
        0,
        0,
        65,
        108,
        101,
        120,
        97,
        110,
        100,
        101,
        114,
        32,
        86,
        114,
        121,
        110,
        105,
        111,
        116,
        105,
        115,
        0,
        0,
        73,
        32,
        108,
        111,
        118,
        101,
        32,
        121,
        111,
        117,
    ]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")

    # test 0
    y_pil = F.adjust_sharpness(x_pil, 1)
    y_np = np.array(y_pil)
    torch.testing.assert_close(y_np, x_np)

    # test 1
    y_pil = F.adjust_sharpness(x_pil, 0.5)
    y_np = np.array(y_pil)
    y_ans = [
        75,
        121,
        114,
        105,
        97,
        107,
        105,
        32,
        66,
        111,
        117,
        114,
        99,
        104,
        97,
        30,
        30,
        74,
        103,
        96,
        114,
        97,
        110,
        100,
        101,
        114,
        32,
        81,
        103,
        108,
        102,
        101,
        107,
        116,
        105,
        115,
        0,
        0,
        73,
        32,
        108,
        111,
        118,
        101,
        32,
        121,
        111,
        117,
    ]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)

    # test 2
    y_pil = F.adjust_sharpness(x_pil, 2)
    y_np = np.array(y_pil)
    y_ans = [
        75,
        121,
        114,
        105,
        97,
        107,
        105,
        32,
        66,
        111,
        117,
        114,
        99,
        104,
        97,
        0,
        0,
        46,
        118,
        111,
        132,
        97,
        110,
        100,
        101,
        114,
        32,
        95,
        135,
        146,
        126,
        112,
        119,
        116,
        105,
        115,
        0,
        0,
        73,
        32,
        108,
        111,
        118,
        101,
        32,
        121,
        111,
        117,
    ]
    y_ans = np.array(y_ans, dtype=np.uint8).reshape(x_shape)
    torch.testing.assert_close(y_np, y_ans)

    # test 3
    x_shape = [2, 2, 3]
    x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
    x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
    x_pil = Image.fromarray(x_np, mode="RGB")
    x_th = torch.tensor(x_np.transpose(2, 0, 1))
    y_pil = F.adjust_sharpness(x_pil, 2)
    y_np = np.array(y_pil).transpose(2, 0, 1)
    y_th = F.adjust_sharpness(x_th, 2)
    torch.testing.assert_close(y_np, y_th.numpy())
//  [python]
def test_model_1b_logits(self):
        input_ids = [[1, 306, 4658, 278, 6593, 310, 2834, 338]]
        model = OlmoForCausalLM.from_pretrained("allenai/OLMo-1B-hf", device_map="auto")
        out = model(torch.tensor(input_ids)).logits
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[2.2869, 0.3315, 0.9876, 1.4146, 1.8804, 2.0430, 1.7055, 1.2065]])
        torch.allclose(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([2.5551, -1.1230, 11.0510, 12.4977, 7.9651, 7.2342, 6.1885, 7.8340, 9.9847, 12.6695, 12.2345, 10.7970, 8.4749, 14.2483, 12.9588, 13.9233, 11.0496, 5.5749, 7.4466, 7.7914, 6.8440, 5.8951, 4.8180, 4.1935, 4.5216, 4.7256, 3.9553, 12.2870, 12.4990, 8.1591])  # fmt: skip
        torch.allclose(out[0, 0, :30], EXPECTED_SLICE, atol=1e-2, rtol=1e-2)
//  [python]
def test_model_1b_logits(self):
        input_ids = [[1, 306, 4658, 278, 6593, 310, 2834, 338]]
        model = OlmoForCausalLM.from_pretrained("allenai/OLMo-1B-hf", device_map="auto")
        out = model(torch.tensor(input_ids)).logits
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[2.2869, 0.3315, 0.9876, 1.4146, 1.8804, 2.0430, 1.7055, 1.2065]])
        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([2.5551, -1.1230, 11.0510, 12.4977, 7.9651, 7.2342, 6.1885, 7.8340, 9.9847, 12.6695, 12.2345, 10.7970, 8.4749, 14.2483, 12.9588, 13.9233, 11.0496, 5.5749, 7.4466, 7.7914, 6.8440, 5.8951, 4.8180, 4.1935, 4.5216, 4.7256, 3.9553, 12.2870, 12.4990, 8.1591])  # fmt: skip
        torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, atol=1e-2, rtol=1e-2)
//  [python]
def test_model_7b_twin_2t_logits(self):
        input_ids = [[1, 306, 4658, 278, 6593, 310, 2834, 338]]
        model = OlmoForCausalLM.from_pretrained("allenai/OLMo-7B-Twin-2T-hf", device_map="auto")
        out = model(torch.tensor(input_ids)).logits
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-0.3636, -0.3825, -0.4800, -0.3696, -0.8388, -0.9737, -0.9849, -0.8356]])
        torch.testing.assert_allclose(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([-2.0833, -1.9234, 8.7312, 7.8049, 1.0372, 0.8941, 3.1548, 1.8502, 5.5511, 5.5793, 8.1166, 4.5906, 1.8691, 11.6377, 8.9858, 11.6447, 7.4549, 1.4725, 2.8399, 2.7568, 1.4011, 1.6958, 0.5572, 0.5231, 0.3068, 0.5364, 0.6769, 7.9636, 8.2379, 1.7950])  # fmt: skip
        torch.testing.assert_allclose(out[0, 0, :30], EXPECTED_SLICE, atol=1e-2, rtol=1e-2)
//  [python]
def test_model_7b_twin_2t_logits(self):
        input_ids = [[1, 306, 4658, 278, 6593, 310, 2834, 338]]
        model = OlmoForCausalLM.from_pretrained("allenai/OLMo-7B-Twin-2T-hf", device_map="auto")
        out = model(torch.tensor(input_ids)).logits
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[-0.3636, -0.3825, -0.4800, -0.3696, -0.8388, -0.9737, -0.9849, -0.8356]])
        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([-2.0833, -1.9234, 8.7312, 7.8049, 1.0372, 0.8941, 3.1548, 1.8502, 5.5511, 5.5793, 8.1166, 4.5906, 1.8691, 11.6377, 8.9858, 11.6447, 7.4549, 1.4725, 2.8399, 2.7568, 1.4011, 1.6958, 0.5572, 0.5231, 0.3068, 0.5364, 0.6769, 7.9636, 8.2379, 1.7950])  # fmt: skip
        torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, atol=1e-2, rtol=1e-2)
//  [python]
def test_model_7b_logits(self):
        input_ids = [[1, 306, 4658, 278, 6593, 310, 2834, 338]]
        model = OlmoForCausalLM.from_pretrained("allenai/OLMo-7B-hf", device_map="auto")
        out = model(torch.tensor(input_ids)).logits
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[0.0271, 0.0249, -0.0578, -0.0870, 0.0167, 0.0710, 0.1002, 0.0677]])
        torch.allclose(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([-1.7433, -1.6685, 7.4941, 6.1506, 0.1364, -0.1127, 1.3224, 4.5458, 4.2068, 5.8296, 7.4723, 2.7925, 3.1245, 10.8872, 10.0758, 10.6717, 7.0945, 1.2398, 3.6766, 4.2365, 2.5655, 2.2222, 1.7418, 0.5223, 0.7753, 1.0938, 0.6723, 6.2522, 6.2264, 1.8105])  # fmt: skip
        torch.allclose(out[0, 0, :30], EXPECTED_SLICE, atol=1e-2, rtol=1e-2)
//  [python]
def test_model_7b_logits(self):
        input_ids = [[1, 306, 4658, 278, 6593, 310, 2834, 338]]
        model = OlmoForCausalLM.from_pretrained("allenai/OLMo-7B-hf", device_map="auto")
        out = model(torch.tensor(input_ids)).logits
        # Expected mean on dim = -1
        EXPECTED_MEAN = torch.tensor([[0.0271, 0.0249, -0.0578, -0.0870, 0.0167, 0.0710, 0.1002, 0.0677]])
        torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
        # slicing logits[0, 0, 0:30]
        EXPECTED_SLICE = torch.tensor([-1.7433, -1.6685, 7.4941, 6.1506, 0.1364, -0.1127, 1.3224, 4.5458, 4.2068, 5.8296, 7.4723, 2.7925, 3.1245, 10.8872, 10.0758, 10.6717, 7.0945, 1.2398, 3.6766, 4.2365, 2.5655, 2.2222, 1.7418, 0.5223, 0.7753, 1.0938, 0.6723, 6.2522, 6.2264, 1.8105])  # fmt: skip
        torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, atol=1e-2, rtol=1e-2)
//  [python]
def test_sobolengine_reset(self, scramble: bool = False):
        ref_sample = self._sobol_reference_samples(scramble=scramble)
        engine = torch.quasirandom.SobolEngine(2, scramble=scramble, seed=123456)
        _ = engine.draw(n=len(ref_sample) // 2)
        engine.reset()
        self.assertEqual(engine.num_generated, 0)
        sample = engine.draw(n=len(ref_sample))
        self.assertEqual(sample, ref_sample)
//  [python]
def test_sobolengine_reset(self, scramble: bool = False):
        ref_sample = self._sobol_reference_samples(scramble=scramble)
        engine = torch.quasirandom.SobolEngine(2, scramble=scramble, seed=123456)
        _ = engine.draw(n=len(ref_sample) // 2)
        engine.reset()
        self.assertEqual(engine.num_generated, 0)
        sample = engine.draw(n=len(ref_sample))
        torch.testing.assert_close(sample, ref_sample)
//  [python]
def test_sobolengine_fast_forward(self, scramble: bool = False):
        ref_sample = self._sobol_reference_samples(scramble=scramble)
        engine = torch.quasirandom.SobolEngine(2, scramble=scramble, seed=123456)
        engine.fast_forward(4)
        sample = engine.draw(n=4)
        self.assertEqual(sample, ref_sample[4:])
        # alternate fast forwarding with sampling
        engine.reset()
        even_draws = []
        for i in range(8):
            if i % 2 == 0:
                even_draws.append(engine.draw())
            else:
                engine.fast_forward(1)
        self.assertEqual(
            ref_sample[[i for i in range(8) if i % 2 == 0]],
            torch.from_numpy(np.concatenate(even_draws)),
        )
//  [python]
def test_sobolengine_fast_forward(self, scramble: bool = False):
        ref_sample = self._sobol_reference_samples(scramble=scramble)
        engine = torch.quasirandom.SobolEngine(2, scramble=scramble, seed=123456)
        engine.fast_forward(4)
        sample = engine.draw(n=4)
        torch.testing.assert_close(sample, ref_sample[4:])
        # alternate fast forwarding with sampling
        engine.reset()
        even_draws = []
        for i in range(8):
            if i % 2 == 0:
                even_draws.append(engine.draw())
            else:
                engine.fast_forward(1)
        torch.testing.assert_close(
            ref_sample[[i for i in range(8) if i % 2 == 0]],
            torch.from_numpy(np.concatenate(even_draws)),
        )
//  [python]
def test_sobolengine_distribution(self, scramble=False):
        d = 50
        engine = torch.quasirandom.SobolEngine(d, scramble=scramble, seed=123456)
        sample = engine.draw(1024)
        self.assertEqual(torch.mean(sample, dim=0), torch.full((d,), 0.5), atol=2, rtol=2)
        self.assertEqual(np.percentile(sample, 25, axis=0), np.repeat(0.25, d), atol=2, rtol=2)
        self.assertEqual(np.percentile(sample, 75, axis=0), np.repeat(0.75, d), atol=2, rtol=2)
//  [python]
def test_sobolengine_distribution(self, scramble=False):
        d = 50
        engine = torch.quasirandom.SobolEngine(d, scramble=scramble, seed=123456)
        sample = engine.draw(1024)
        torch.testing.assert_close(
            torch.mean(sample, dim=0), torch.full((d,), 0.5), atol=2, rtol=2
        )
        torch.testing.assert_close(
            np.percentile(sample, 25, axis=0), np.repeat(0.25, d), atol=2, rtol=2
        )
        torch.testing.assert_close(
            np.percentile(sample, 75, axis=0), np.repeat(0.75, d), atol=2, rtol=2
        )
//  [python]
def test_lower_tail():
    t.assert_close(
        lower_tail(
            Normal(0.0, 1.0),
            0.5,
        ),
        torch.tensor([-0.6745]),
    )

    t.assert_close(
        lower_tail(
            Uniform(0.0, 1.0),
            0.5,
        ),
        torch.tensor([0.25]),
    )
//  [python]
def test_lower_tail():
    torch.testing.assert_close(
        lower_tail(
            Normal(0.0, 1.0),
            0.5,
        ),
        torch.tensor([-0.6745]),
    )

    torch.testing.assert_close(
        lower_tail(
            Uniform(0.0, 1.0),
            0.5,
        ),
        torch.tensor([0.25]),
    )
//  [python]
def _test_matrix_root_inverse_residuals(
        self,
        A: torch.Tensor,
        X_hat: torch.Tensor,
        root: int,
        exponent_multiplier: float,
        expected_relative_error: torch.Tensor,
        expected_relative_residual: torch.Tensor,
    ) -> None:
        (
            actual_relative_error,
            actual_relative_residual,
        ) = compute_matrix_root_inverse_residuals(
            A=A,
            X_hat=X_hat,
            root=root,
            epsilon=0.0,
            exponent_multiplier=exponent_multiplier,
        )
        t.assert_close(
            actual_relative_error,
            expected_relative_error,
        )
        t.assert_close(
            actual_relative_residual,
            expected_relative_residual,
        )
//  [python]
def _test_matrix_root_inverse_residuals(
        self,
        A: torch.Tensor,
        X_hat: torch.Tensor,
        root: int,
        exponent_multiplier: float,
        expected_relative_error: torch.Tensor,
        expected_relative_residual: torch.Tensor,
    ) -> None:
        (
            actual_relative_error,
            actual_relative_residual,
        ) = compute_matrix_root_inverse_residuals(
            A=A,
            X_hat=X_hat,
            root=root,
            epsilon=0.0,
            exponent_multiplier=exponent_multiplier,
        )
        torch.testing.assert_close(
            actual_relative_error,
            expected_relative_error,
        )
        torch.testing.assert_close(
            actual_relative_residual,
            expected_relative_residual,
        )
//  [python]
def test_matrix_inverse_root(self) -> None:
        A = torch.tensor([[1.0, 0.0], [0.0, 4.0]])
        root = 2
        exponent_multiplier = 1.0
        actual_inverse_root = torch.tensor([[1.0, 0.0], [0.0, 0.5]])
        with self.subTest("Test with diagonal case."):
            self.assertEqual(
                actual_inverse_root,
                matrix_inverse_root(
                    A,
                    root=root,
                    exponent_multiplier=exponent_multiplier,
                    is_diagonal=True,
                ),
            )
        with self.subTest("Test with EIGEN."):
            self.assertEqual(
                actual_inverse_root,
                matrix_inverse_root(
                    A,
                    root=root,
                    exponent_multiplier=exponent_multiplier,
                    root_inv_method=RootInvMethod.EIGEN,
                    is_diagonal=False,
                ),
            )
        with self.subTest("Test with NEWTON."):
            self.assertEqual(
                actual_inverse_root,
                matrix_inverse_root(
                    A,
                    root=root,
                    exponent_multiplier=exponent_multiplier,
                    root_inv_method=RootInvMethod.NEWTON,
                    is_diagonal=False,
                ),
            )
//  [python]
def test_matrix_inverse_root(self) -> None:
        A = torch.tensor([[1.0, 0.0], [0.0, 4.0]])
        root = 2
        exponent_multiplier = 1.0
        actual_inverse_root = torch.tensor([[1.0, 0.0], [0.0, 0.5]])
        with self.subTest("Test with diagonal case."):
            torch.testing.assert_close(
                actual_inverse_root,
                matrix_inverse_root(
                    A,
                    root=root,
                    exponent_multiplier=exponent_multiplier,
                    is_diagonal=True,
                ),
            )
        with self.subTest("Test with EIGEN."):
            torch.testing.assert_close(
                actual_inverse_root,
                matrix_inverse_root(
                    A,
                    root=root,
                    exponent_multiplier=exponent_multiplier,
                    root_inv_method=RootInvMethod.EIGEN,
                    is_diagonal=False,
                ),
            )
        with self.subTest("Test with NEWTON."):
            torch.testing.assert_close(
                actual_inverse_root,
                matrix_inverse_root(
                    A,
                    root=root,
                    exponent_multiplier=exponent_multiplier,
                    root_inv_method=RootInvMethod.NEWTON,
                    is_diagonal=False,
                ),
            )
//  [python]
def test_calib_and_quantize(self, backend: Backend) -> None:
        model = SimpleModule()
        quantizer = Quantizer(backend=backend)
        dummy_input = torch.randn((1, 2, 5, 5))
        original_output = model(dummy_input)

        calib_model = quantizer.calib(model)
        calib_output = calib_model(dummy_input)
        self.assertTrue(torch.equal(original_output, calib_output))

        quant_model = quantizer.quantize(model)
        quant_output = quant_model(dummy_input)
        self.assertFalse(torch.equal(original_output, quant_output))
        torch.testing.assert_allclose(quant_output, original_output, rtol=0.1, atol=0.5)

        if backend == Backend.MNN:
            _, compress_proto = quantizer.export_mnn_params(quant_model, dummy_input)
            self.assertEqual(len(compress_proto.algo[0].quant_params.layer), 3)
//  [python]
def test_calib_and_quantize(self, backend: Backend) -> None:
        model = SimpleModule()
        quantizer = Quantizer(backend=backend)
        dummy_input = torch.randn((1, 2, 5, 5))
        original_output = model(dummy_input)

        calib_model = quantizer.calib(model)
        calib_output = calib_model(dummy_input)
        self.assertTrue(torch.equal(original_output, calib_output))

        quant_model = quantizer.quantize(model)
        quant_output = quant_model(dummy_input)
        self.assertFalse(torch.equal(original_output, quant_output))
        torch.testing.assert_close(quant_output, original_output, rtol=0.1, atol=0.5)

        if backend == Backend.MNN:
            _, compress_proto = quantizer.export_mnn_params(quant_model, dummy_input)
            self.assertEqual(len(compress_proto.algo[0].quant_params.layer), 3)