jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_mixup(self):
        class Model(torch.nn.Module):
            def forward(self, batch):
                return batch

        images = torch.stack(
            [
                torch.ones(1, 12, 12),
                torch.ones(1, 12, 12) * 2,
                torch.ones(1, 12, 12) * 3,
            ]
        ).to(torch.float)
        labels = torch.tensor([1, 2, 3])

        class args:
            pass

        args = args()
        args.mixup_alpha = 0.5
        args.precision = "32.32"

        model = AugmentationModel(Model(), use_mixup=True, use_cutmix=False, args=args)

        def dummy_sample(bs=3):
            return torch.tensor([0.5, 0.5, 0.5])

        model._sample_mixup_coefficients = dummy_sample
        model = poptorch.inferenceModel(model)
        mixed_images, all_coeffs = model(images)

        correct_mixed_images = torch.stack(
            [
                torch.ones(1, 12, 12) * 2,  # 0.5 * 1 + 0.5 * 3
                torch.ones(1, 12, 12) * 1.5,  # 0.5 * 2 + 0.5 * 1
                torch.ones(1, 12, 12) * 2.5,  # 0.5 * 3 + 0.5 * 2
            ]
        ).to(torch.float)

        torch.testing.assert_allclose(mixed_images, correct_mixed_images)
        torch.testing.assert_allclose(all_coeffs[0], dummy_sample())

        all_labels, weights = model.mix_labels(labels, all_coeffs)
        torch.testing.assert_allclose(
            actual=all_labels[0],
            expected=labels,
            rtol=0,
            atol=0,
        )
        torch.testing.assert_allclose(
            actual=all_labels[1],
            expected=torch.tensor([3, 1, 2]),
            rtol=0,
            atol=0,
        )
        torch.testing.assert_allclose(weights[0], dummy_sample())
        torch.testing.assert_allclose(weights[1], 1.0 - dummy_sample())
//  [python]
def test_mixup(self):
        class Model(torch.nn.Module):
            def forward(self, batch):
                return batch

        images = torch.stack(
            [
                torch.ones(1, 12, 12),
                torch.ones(1, 12, 12) * 2,
                torch.ones(1, 12, 12) * 3,
            ]
        ).to(torch.float)
        labels = torch.tensor([1, 2, 3])

        class args:
            pass

        args = args()
        args.mixup_alpha = 0.5
        args.precision = "32.32"

        model = AugmentationModel(Model(), use_mixup=True, use_cutmix=False, args=args)

        def dummy_sample(bs=3):
            return torch.tensor([0.5, 0.5, 0.5])

        model._sample_mixup_coefficients = dummy_sample
        model = poptorch.inferenceModel(model)
        mixed_images, all_coeffs = model(images)

        correct_mixed_images = torch.stack(
            [
                torch.ones(1, 12, 12) * 2,  # 0.5 * 1 + 0.5 * 3
                torch.ones(1, 12, 12) * 1.5,  # 0.5 * 2 + 0.5 * 1
                torch.ones(1, 12, 12) * 2.5,  # 0.5 * 3 + 0.5 * 2
            ]
        ).to(torch.float)

        torch.testing.assert_close(mixed_images, correct_mixed_images)
        torch.testing.assert_close(all_coeffs[0], dummy_sample())

        all_labels, weights = model.mix_labels(labels, all_coeffs)
        torch.testing.assert_close(
            actual=all_labels[0],
            expected=labels,
            rtol=0,
            atol=0,
        )
        torch.testing.assert_close(
            actual=all_labels[1],
            expected=torch.tensor([3, 1, 2]),
            rtol=0,
            atol=0,
        )
        torch.testing.assert_close(weights[0], dummy_sample())
        torch.testing.assert_close(weights[1], 1.0 - dummy_sample())
//  [python]
def test_type_as_cat(self):
        with inline_fusion_groups():
            def eager(x, y):
                return torch.cat((x, y.type_as(x)), dim=1)
            for dtype1, dtype2 in product(self.dtypes, self.dtypes):
                x = torch.randint(2, (1, 13,)).to(dtype1)
                zero = torch.tensor([[0]]).to(dtype2)
                one = torch.tensor([[1]]).to(dtype2)
                script = torch.jit.trace(eager, (x, zero))
                for _ in range(3):
                    torch.testing.assert_allclose(
                        script(x, zero),
                        eager(x, zero))
                    torch.testing.assert_allclose(
                        script(x, one),
                        eager(x, one))
                self.assertAllFused(script.graph_for(x, one))
//  [python]
def test_type_as_cat(self):
        with inline_fusion_groups():
            def eager(x, y):
                return torch.cat((x, y.type_as(x)), dim=1)
            for dtype1, dtype2 in product(self.dtypes, self.dtypes):
                x = torch.randint(2, (1, 13,)).to(dtype1)
                zero = torch.tensor([[0]]).to(dtype2)
                one = torch.tensor([[1]]).to(dtype2)
                script = torch.jit.trace(eager, (x, zero))
                for _ in range(3):
                    torch.testing.assert_close(
                        script(x, zero),
                        eager(x, zero))
                    torch.testing.assert_close(
                        script(x, one),
                        eager(x, one))
                self.assertAllFused(script.graph_for(x, one))
//  [python]
def _test_fwd_bwd(self, fn):
        x = torch.arange(-10, 10, dtype=torch.float32, requires_grad=True)
        xs = torch.arange(-10, 10, dtype=torch.float32, requires_grad=True)
        script = torch.jit.script(fn)
        for i in range(11):
            y = fn(x)
            g0 = torch.rand_like(y)
            y.backward(g0)

            ys = script(xs)
            ys.backward(g0)

            with torch.no_grad():
                x -= 0.1 * x.grad
                xs -= 0.1 * xs.grad
                x.grad = None
                xs.grad = None
        torch.testing.assert_allclose(y, ys)
//  [python]
def _test_fwd_bwd(self, fn):
        x = torch.arange(-10, 10, dtype=torch.float32, requires_grad=True)
        xs = torch.arange(-10, 10, dtype=torch.float32, requires_grad=True)
        script = torch.jit.script(fn)
        for i in range(11):
            y = fn(x)
            g0 = torch.rand_like(y)
            y.backward(g0)

            ys = script(xs)
            ys.backward(g0)

            with torch.no_grad():
                x -= 0.1 * x.grad
                xs -= 0.1 * xs.grad
                x.grad = None
                xs.grad = None
        torch.testing.assert_close(y, ys)
//  [python]
def test_unrolled_cat(self):
        with inline_fusion_groups():
            def eager(x):
                ret = torch.empty(0)
                for i in range(x.shape[0]):
                    ret = torch.cat([ret, x[i].relu()])
                return ret
            script = torch.jit.script(eager)

            # Warm up with size=1 tensor; since the loop iterates once the
            # profile data will be "burned in" assuming size=1, and then
            # unrolled.
            x = torch.ones(1, 1)
            for _ in range(3):
                script(x)

            torch.testing.assert_allclose(eager(x), script(x))

            # Now when an input hits the unrolled path, it will produce an
            # incorrectly-sized tensor, since size=1 has been burned in.
            x = torch.ones((8, 1))
            torch.testing.assert_allclose(eager(x), script(x))
//  [python]
def test_unrolled_cat(self):
        with inline_fusion_groups():
            def eager(x):
                ret = torch.empty(0)
                for i in range(x.shape[0]):
                    ret = torch.cat([ret, x[i].relu()])
                return ret
            script = torch.jit.script(eager)

            # Warm up with size=1 tensor; since the loop iterates once the
            # profile data will be "burned in" assuming size=1, and then
            # unrolled.
            x = torch.ones(1, 1)
            for _ in range(3):
                script(x)

            torch.testing.assert_close(eager(x), script(x))

            # Now when an input hits the unrolled path, it will produce an
            # incorrectly-sized tensor, since size=1 has been burned in.
            x = torch.ones((8, 1))
            torch.testing.assert_close(eager(x), script(x))
//  [python]
def _test_binary_f1_score_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        threshold: float = 0.5,
    ) -> None:
        input_np = input.numpy().round()
        target_np = target.squeeze().numpy()
        sklearn_result = torch.tensor(f1_score(target_np, input_np)).to(torch.float32)

        torch.testing.assert_allclose(
            my_binary_f1_score(input, target, threshold=threshold),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_binary_f1_score_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        threshold: float = 0.5,
    ) -> None:
        input_np = input.numpy().round()
        target_np = target.squeeze().numpy()
        sklearn_result = torch.tensor(f1_score(target_np, input_np)).to(torch.float32)

        torch.testing.assert_close(
            my_binary_f1_score(input, target, threshold=threshold),
            sklearn_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_masked_fill(self):
        dtypes = [
            torch.int8,
            torch.int16,
            torch.int32,
            torch.int64,
            # TODO: Add back when https://github.com/pytorch/pytorch/issues/55905 is closed
            # torch.float16,
            torch.float32,
            torch.float64,
            torch.bool,
        ]
        sizes = [(2,), (4, 4)]
        for self_dtype, device, scalar_val, size in product(dtypes, self.devices, [0.4, 3], sizes):
            input_v = self.data_for(self_dtype, device, size=size)
            mask = self.data_for(torch.bool, device, size=size)

            def fn(input_v, mask):
                return torch.masked_fill(input_v, mask, scalar_val)
            ref = fn(input_v, mask)
            try:
                t = torch.jit.trace(fn, (input_v, mask))
                torch.testing.assert_allclose(ref, t(input_v, mask))
                print(torch.jit.last_executed_optimized_graph())
                self.assertLastGraphAllFused()
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(self_dtype), op.__name__, device, str(size)])
                )
//  [python]
def test_masked_fill(self):
        dtypes = [
            torch.int8,
            torch.int16,
            torch.int32,
            torch.int64,
            # TODO: Add back when https://github.com/pytorch/pytorch/issues/55905 is closed
            # torch.float16,
            torch.float32,
            torch.float64,
            torch.bool,
        ]
        sizes = [(2,), (4, 4)]
        for self_dtype, device, scalar_val, size in product(dtypes, self.devices, [0.4, 3], sizes):
            input_v = self.data_for(self_dtype, device, size=size)
            mask = self.data_for(torch.bool, device, size=size)

            def fn(input_v, mask):
                return torch.masked_fill(input_v, mask, scalar_val)
            ref = fn(input_v, mask)
            try:
                t = torch.jit.trace(fn, (input_v, mask))
                torch.testing.assert_close(ref, t(input_v, mask))
                print(torch.jit.last_executed_optimized_graph())
                self.assertLastGraphAllFused()
            except Exception as e:
                raise RuntimeError(
                    " ".join(["Failed:", str(self_dtype), op.__name__, device, str(size)])
                )
//  [python]
def skipping_quantization_example():
    class MyModel(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
            self.conv2 = SkipQuantization(nn.Conv2d(8, 8, kernel_size=3, padding=1))  # can use the wrapper to skip

        def forward(self, x):
            return self.conv2(self.conv1(x))

    module = MyModel()

    # Initialize the quantization utility, register layers to skip, and quantize the module
    q_util = SelectiveQuantizer()
    q_util.register_skip_quantization(layer_names={"conv1"})  # can also configure skip by layer names
    q_util.quantize_module(module)

    print(module)  # You should expect to see Conv2d

    x = torch.rand(1, 3, 32, 32)
    with torch.no_grad():
        y = module(x)
        t.assert_close(y.size(), (1, 8, 32, 32))
//  [python]
def skipping_quantization_example():
    class MyModel(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
            self.conv2 = SkipQuantization(nn.Conv2d(8, 8, kernel_size=3, padding=1))  # can use the wrapper to skip

        def forward(self, x):
            return self.conv2(self.conv1(x))

    module = MyModel()

    # Initialize the quantization utility, register layers to skip, and quantize the module
    q_util = SelectiveQuantizer()
    q_util.register_skip_quantization(layer_names={"conv1"})  # can also configure skip by layer names
    q_util.quantize_module(module)

    print(module)  # You should expect to see Conv2d

    x = torch.rand(1, 3, 32, 32)
    with torch.no_grad():
        y = module(x)
        torch.testing.assert_close(y.size(), (1, 8, 32, 32))
//  [python]
def test_cutmix(self):
        class Model(torch.nn.Module):
            def forward(self, batch):
                return batch

        images = torch.stack(
            [
                torch.ones(1, 10, 10),
                torch.ones(1, 10, 10) * 2,
                torch.ones(1, 10, 10) * 3,
            ]
        ).to(torch.float)
        labels = torch.tensor([1, 2, 3])

        class args:
            pass

        args = args()
        # 0.75 should result in 5x5 cut boxes.
        args.cutmix_lambda_low = 0.75
        args.cutmix_lambda_high = 0.75
        args.cutmix_disable_prob = 0.0
        args.precision = "32.32"

        model = poptorch.inferenceModel(
            AugmentationModel(
                Model(),
                use_mixup=False,
                use_cutmix=True,
                args=args,
            )
        )
        cutmixed_images, all_coeffs = model(images)

        torch.testing.assert_allclose(all_coeffs[0], torch.tensor(0.75))
        torch.testing.assert_allclose(
            torch.sum(cutmixed_images, dim=(1, 2, 3)),
            torch.tensor(
                [
                    0.75 * images[0].sum() + 0.25 * images[1].sum(),
                    0.75 * images[1].sum() + 0.25 * images[2].sum(),
                    0.75 * images[2].sum() + 0.25 * images[0].sum(),
                ]
            ),
        )
        torch.testing.assert_allclose(
            actual=torch.unique(cutmixed_images[0], sorted=True),
            expected=torch.tensor([1.0, 2.0]),
            rtol=0,
            atol=0,
        )
        torch.testing.assert_allclose(
            actual=torch.unique(cutmixed_images[1], sorted=True),
            expected=torch.tensor([2.0, 3.0]),
            rtol=0,
            atol=0,
        )
        torch.testing.assert_allclose(
            actual=torch.unique(cutmixed_images[2], sorted=True),
            expected=torch.tensor([1.0, 3.0]),
            rtol=0,
            atol=0,
        )

        all_labels, weights = model.mix_labels(labels, all_coeffs)
        torch.testing.assert_allclose(
            actual=all_labels[0],
            expected=labels,
            rtol=0,
            atol=0,
        )
        torch.testing.assert_allclose(
            actual=all_labels[1],
            expected=torch.tensor([2, 3, 1]),
            rtol=0,
            atol=0,
        )
        torch.testing.assert_allclose(weights[0], torch.tensor(0.75))
        torch.testing.assert_allclose(weights[1], 1.0 - torch.tensor(0.75))
//  [python]
def test_cutmix(self):
        class Model(torch.nn.Module):
            def forward(self, batch):
                return batch

        images = torch.stack(
            [
                torch.ones(1, 10, 10),
                torch.ones(1, 10, 10) * 2,
                torch.ones(1, 10, 10) * 3,
            ]
        ).to(torch.float)
        labels = torch.tensor([1, 2, 3])

        class args:
            pass

        args = args()
        # 0.75 should result in 5x5 cut boxes.
        args.cutmix_lambda_low = 0.75
        args.cutmix_lambda_high = 0.75
        args.cutmix_disable_prob = 0.0
        args.precision = "32.32"

        model = poptorch.inferenceModel(
            AugmentationModel(
                Model(),
                use_mixup=False,
                use_cutmix=True,
                args=args,
            )
        )
        cutmixed_images, all_coeffs = model(images)

        torch.testing.assert_close(all_coeffs[0], torch.tensor(0.75))
        torch.testing.assert_close(
            torch.sum(cutmixed_images, dim=(1, 2, 3)),
            torch.tensor(
                [
                    0.75 * images[0].sum() + 0.25 * images[1].sum(),
                    0.75 * images[1].sum() + 0.25 * images[2].sum(),
                    0.75 * images[2].sum() + 0.25 * images[0].sum(),
                ]
            ),
        )
        torch.testing.assert_close(
            actual=torch.unique(cutmixed_images[0], sorted=True),
            expected=torch.tensor([1.0, 2.0]),
            rtol=0,
            atol=0,
        )
        torch.testing.assert_close(
            actual=torch.unique(cutmixed_images[1], sorted=True),
            expected=torch.tensor([2.0, 3.0]),
            rtol=0,
            atol=0,
        )
        torch.testing.assert_close(
            actual=torch.unique(cutmixed_images[2], sorted=True),
            expected=torch.tensor([1.0, 3.0]),
            rtol=0,
            atol=0,
        )

        all_labels, weights = model.mix_labels(labels, all_coeffs)
        torch.testing.assert_close(
            actual=all_labels[0],
            expected=labels,
            rtol=0,
            atol=0,
        )
        torch.testing.assert_close(
            actual=all_labels[1],
            expected=torch.tensor([2, 3, 1]),
            rtol=0,
            atol=0,
        )
        torch.testing.assert_close(weights[0], torch.tensor(0.75))
        torch.testing.assert_close(weights[1], 1.0 - torch.tensor(0.75))
//  [python]
def test_binary_f1_score_thresholding(self) -> None:
        num_classes = 2
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))

        # test threshold larger than every prediction gives 0
        torch.testing.assert_allclose(
            my_binary_f1_score(input, target, threshold=2),
            torch.Tensor([0.0]).sum(),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_binary_f1_score_thresholding(self) -> None:
        num_classes = 2
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))

        # test threshold larger than every prediction gives 0
        torch.testing.assert_close(
            my_binary_f1_score(input, target, threshold=2),
            torch.Tensor([0.0]).sum(),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_log_ndtr():
    rng = numpy.random.default_rng(0xDEADBEEF)

    x = rng.random((32,), dtype=numpy.float64)

    actual = log_ndtr(torch.tensor(x, dtype=torch.float))

    assert torch.isfinite(actual).all()

    torch.testing.assert_allclose(
        actual,
        torch.tensor(scipy.special.log_ndtr(x), dtype=torch.float),
    )
//  [python]
def test_log_ndtr():
    rng = numpy.random.default_rng(0xDEADBEEF)

    x = rng.random((32,), dtype=numpy.float64)

    actual = log_ndtr(torch.tensor(x, dtype=torch.float))

    assert torch.isfinite(actual).all()

    torch.testing.assert_close(
        actual,
        torch.tensor(scipy.special.log_ndtr(x), dtype=torch.float),
    )
//  [python]
def test_binary_f1_score_with_0s(self) -> None:
        num_classes = 2

        # test input all 0s
        input = torch.zeros(size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))

        t.assert_close(
            my_binary_f1_score(input, target),
            torch.Tensor([0.0]).sum(),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # test target all 0s
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.zeros(size=(BATCH_SIZE,))

        t.assert_close(
            my_binary_f1_score(input, target),
            torch.Tensor([0.0]).sum(),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_binary_f1_score_with_0s(self) -> None:
        num_classes = 2

        # test input all 0s
        input = torch.zeros(size=(BATCH_SIZE,))
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))

        torch.testing.assert_close(
            my_binary_f1_score(input, target),
            torch.Tensor([0.0]).sum(),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # test target all 0s
        input = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        target = torch.zeros(size=(BATCH_SIZE,))

        torch.testing.assert_close(
            my_binary_f1_score(input, target),
            torch.Tensor([0.0]).sum(),
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_video_mixed_cut_from_padding(video_cut):

    orig_video, orig_audio = video_cut.load_video()

    video_cut = video_cut.pad(duration=10.0)
    assert video_cut.has_video
    assert video_cut.video.fps == FPS
    assert video_cut.video.width == WIDTH
    assert video_cut.video.height == HEIGHT
    assert video_cut.video.num_frames == 25 * 10

    # Load all audio channels with video
    video, audio = video_cut.load_video()
    assert video.dtype == torch.uint8
    assert video.shape == (25 * 10, COLOR, HEIGHT, WIDTH)
    assert audio.dtype == torch.float32
    assert audio.shape == (AUDIO_CHANNELS, 48000 * 10)

    t.assert_close(video[:132], orig_video)
    assert (video[132:] == 0).all()
    t.assert_close(audio[:, :253440], orig_audio)
    assert (audio[:, 253440:] == 0).all()
//  [python]
def test_video_mixed_cut_from_padding(video_cut):

    orig_video, orig_audio = video_cut.load_video()

    video_cut = video_cut.pad(duration=10.0)
    assert video_cut.has_video
    assert video_cut.video.fps == FPS
    assert video_cut.video.width == WIDTH
    assert video_cut.video.height == HEIGHT
    assert video_cut.video.num_frames == 25 * 10

    # Load all audio channels with video
    video, audio = video_cut.load_video()
    assert video.dtype == torch.uint8
    assert video.shape == (25 * 10, COLOR, HEIGHT, WIDTH)
    assert audio.dtype == torch.float32
    assert audio.shape == (AUDIO_CHANNELS, 48000 * 10)

    torch.testing.assert_close(video[:132], orig_video)
    assert (video[132:] == 0).all()
    torch.testing.assert_close(audio[:, :253440], orig_audio)
    assert (audio[:, 253440:] == 0).all()
//  [python]
def test_video_mixed_cut_from_appending(video_cut):

    orig_video, orig_audio = video_cut.load_video()

    video_cut = video_cut.append(video_cut)
    assert video_cut.has_video
    assert video_cut.video.fps == FPS
    assert video_cut.video.width == WIDTH
    assert video_cut.video.height == HEIGHT
    assert video_cut.video.num_frames == 132 * 2

    # Load all audio channels with video
    video, audio = video_cut.load_video()
    assert video.dtype == torch.uint8
    assert video.shape == (132 * 2, COLOR, HEIGHT, WIDTH)
    assert audio.dtype == torch.float32
    assert audio.shape == (AUDIO_CHANNELS, 253440 * 2)

    t.assert_close(video[:132], orig_video)
    t.assert_close(video[132:], orig_video)
    t.assert_close(audio[:, :253440], orig_audio)
    t.assert_close(audio[:, 253440:], orig_audio)
//  [python]
def test_video_mixed_cut_from_appending(video_cut):

    orig_video, orig_audio = video_cut.load_video()

    video_cut = video_cut.append(video_cut)
    assert video_cut.has_video
    assert video_cut.video.fps == FPS
    assert video_cut.video.width == WIDTH
    assert video_cut.video.height == HEIGHT
    assert video_cut.video.num_frames == 132 * 2

    # Load all audio channels with video
    video, audio = video_cut.load_video()
    assert video.dtype == torch.uint8
    assert video.shape == (132 * 2, COLOR, HEIGHT, WIDTH)
    assert audio.dtype == torch.float32
    assert audio.shape == (AUDIO_CHANNELS, 253440 * 2)

    torch.testing.assert_close(video[:132], orig_video)
    torch.testing.assert_close(video[132:], orig_video)
    torch.testing.assert_close(audio[:, :253440], orig_audio)
    torch.testing.assert_close(audio[:, 253440:], orig_audio)
//  [python]
def assert_expected(
    actual: Any,
    expected: Any,
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
    check_device=True,
):
    t.assert_close(
        actual,
        expected,
        rtol=rtol,
        atol=atol,
        check_device=check_device,
        msg=f"actual: {actual}, expected: {expected}",
    )
//  [python]
def assert_expected(
    actual: Any,
    expected: Any,
    rtol: Optional[float] = None,
    atol: Optional[float] = None,
    check_device=True,
):
    torch.testing.assert_close(
        actual,
        expected,
        rtol=rtol,
        atol=atol,
        check_device=check_device,
        msg=f"actual: {actual}, expected: {expected}",
    )
//  [python]
def _test_f1_score_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: Optional[int] = None,
        average: Optional[str] = "micro",
    ) -> None:
        if num_classes is None:
            if average == "micro":
                my_compute_result = my_f1_score(input, target)
            else:
                my_compute_result = my_f1_score(input, target, average=average)
        else:
            if average == "micro":
                my_compute_result = my_f1_score(input, target, num_classes=num_classes)
            else:
                my_compute_result = my_f1_score(
                    input, target, num_classes=num_classes, average=average
                )
        target = target.squeeze().flatten()
        input_squeezed = input.squeeze()
        input_label_ids = (
            torch.argmax(input_squeezed, dim=1)
            if input_squeezed.ndim == 2
            else input_squeezed
        )
        input = input_label_ids.flatten()
        compute_result = torch.tensor(f1_score(target, input, average=average)).to(
            torch.float32
        )
        if my_compute_result.shape != compute_result.shape:
            compute_result = torch.from_numpy(
                np.append(compute_result.numpy(), 0.0)
            ).to(torch.float32)
        t.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_f1_score_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: Optional[int] = None,
        average: Optional[str] = "micro",
    ) -> None:
        if num_classes is None:
            if average == "micro":
                my_compute_result = my_f1_score(input, target)
            else:
                my_compute_result = my_f1_score(input, target, average=average)
        else:
            if average == "micro":
                my_compute_result = my_f1_score(input, target, num_classes=num_classes)
            else:
                my_compute_result = my_f1_score(
                    input, target, num_classes=num_classes, average=average
                )
        target = target.squeeze().flatten()
        input_squeezed = input.squeeze()
        input_label_ids = (
            torch.argmax(input_squeezed, dim=1)
            if input_squeezed.ndim == 2
            else input_squeezed
        )
        input = input_label_ids.flatten()
        compute_result = torch.tensor(f1_score(target, input, average=average)).to(
            torch.float32
        )
        if my_compute_result.shape != compute_result.shape:
            compute_result = torch.from_numpy(
                np.append(compute_result.numpy(), 0.0)
            ).to(torch.float32)
        torch.testing.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def validate_mixup_batch(x, y, indices, x_mix, y_perm, mixing):
    # Explicitly check that the batches and labels have been mixed correctly.
    for i in range(x.size(0)):  # Grab N
        j = indices[i]
        # Check the input data
        x_mix_test = (1 - mixing) * x[i] + mixing * x[j]
        t.assert_close(x_mix_test, x_mix[i])
        # Check the label
        perm_label = y[j]
        t.assert_close(perm_label, y_perm[i])
//  [python]
def validate_mixup_batch(x, y, indices, x_mix, y_perm, mixing):
    # Explicitly check that the batches and labels have been mixed correctly.
    for i in range(x.size(0)):  # Grab N
        j = indices[i]
        # Check the input data
        x_mix_test = (1 - mixing) * x[i] + mixing * x[j]
        torch.testing.assert_close(x_mix_test, x_mix[i])
        # Check the label
        perm_label = y[j]
        torch.testing.assert_close(perm_label, y_perm[i])
//  [python]
def _test_binary_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_tasks: int = 1,
        compute_result: Optional[torch.Tensor] = None,
    ) -> None:

        device = "cpu"
        if torch.cuda.is_available():
            device = "cuda"

        if compute_result is None:
            compute_result = self._get_sklearn_equivalent(input, target, device)

        # Get torcheval compute result
        te_compute_result = binary_auprc(
            input.to(device=device),
            target.to(device=device),
            num_tasks=num_tasks,
        )

        torch.testing.assert_allclose(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_binary_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_tasks: int = 1,
        compute_result: Optional[torch.Tensor] = None,
    ) -> None:

        device = "cpu"
        if torch.cuda.is_available():
            device = "cuda"

        if compute_result is None:
            compute_result = self._get_sklearn_equivalent(input, target, device)

        # Get torcheval compute result
        te_compute_result = binary_auprc(
            input.to(device=device),
            target.to(device=device),
            num_tasks=num_tasks,
        )

        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_lora_fsdp_wrap(self):
        with torch.device("meta"):
            model = lora_llama2(
                lora_attn_modules=["q_proj", "v_proj"],
                vocab_size=VOCAB_SIZE,
                num_layers=N_LAYERS,
                num_heads=NUM_HEADS,
                num_kv_heads=NUM_KV_HEADS,
                embed_dim=EMBED_DIM,
                max_seq_len=MAX_SEQ_LEN,
                lora_rank=4,
                lora_alpha=1.0,
            )

        adapter_params = get_adapter_params(model)
        set_trainable_params(model, adapter_params)
        num_lora_ab, num_transformer_layers = _get_n_lora_and_tformer_layers(model)
        with single_box_init():
            lora_wrap_policy = utils.lora_fsdp_wrap_policy(
                modules_to_wrap={TransformerDecoderLayer}
            )
            utils.prepare_model_for_fsdp_with_meta_device(model)
            wrapped_lora = FSDP(
                model,
                auto_wrap_policy=lora_wrap_policy,
                device_id=torch.device("cpu"),
            )

            # After FSDP wrap, nothing should be left on meta device, and LoRA params
            # should be initialized.
            for p in chain(wrapped_lora.parameters(), wrapped_lora.buffers()):
                assert not p.is_meta

            for m in wrapped_lora.modules():
                if isinstance(m, LoRALinear):
                    torch.testing.assert_allclose(
                        m.lora_b.weight, torch.zeros_like(m.lora_b.weight)
                    )

            # Total # FSDP modules should be num_transformer + num_lora_ab + 1
            total_fsdp_submodules = len([m for m in FSDP.fsdp_modules(wrapped_lora)])
            assert total_fsdp_submodules == (num_lora_ab + num_transformer_layers + 1)
            # LoRA a & b linears should be individually wrapped.
            # And TransformerDecoderLayers should be individually wrapped.
            for fsdp_submodule in FSDP.fsdp_modules(wrapped_lora):
                if isinstance(fsdp_submodule.module, nn.Linear):
                    num_lora_ab -= 1
                elif isinstance(fsdp_submodule.module, TransformerDecoderLayer):
                    num_transformer_layers -= 1
            assert num_lora_ab == 0
            assert num_transformer_layers == 0
//  [python]
def test_lora_fsdp_wrap(self):
        with torch.device("meta"):
            model = lora_llama2(
                lora_attn_modules=["q_proj", "v_proj"],
                vocab_size=VOCAB_SIZE,
                num_layers=N_LAYERS,
                num_heads=NUM_HEADS,
                num_kv_heads=NUM_KV_HEADS,
                embed_dim=EMBED_DIM,
                max_seq_len=MAX_SEQ_LEN,
                lora_rank=4,
                lora_alpha=1.0,
            )

        adapter_params = get_adapter_params(model)
        set_trainable_params(model, adapter_params)
        num_lora_ab, num_transformer_layers = _get_n_lora_and_tformer_layers(model)
        with single_box_init():
            lora_wrap_policy = utils.lora_fsdp_wrap_policy(
                modules_to_wrap={TransformerDecoderLayer}
            )
            utils.prepare_model_for_fsdp_with_meta_device(model)
            wrapped_lora = FSDP(
                model,
                auto_wrap_policy=lora_wrap_policy,
                device_id=torch.device("cpu"),
            )

            # After FSDP wrap, nothing should be left on meta device, and LoRA params
            # should be initialized.
            for p in chain(wrapped_lora.parameters(), wrapped_lora.buffers()):
                assert not p.is_meta

            for m in wrapped_lora.modules():
                if isinstance(m, LoRALinear):
                    torch.testing.assert_close(
                        m.lora_b.weight, torch.zeros_like(m.lora_b.weight)
                    )

            # Total # FSDP modules should be num_transformer + num_lora_ab + 1
            total_fsdp_submodules = len([m for m in FSDP.fsdp_modules(wrapped_lora)])
            assert total_fsdp_submodules == (num_lora_ab + num_transformer_layers + 1)
            # LoRA a & b linears should be individually wrapped.
            # And TransformerDecoderLayers should be individually wrapped.
            for fsdp_submodule in FSDP.fsdp_modules(wrapped_lora):
                if isinstance(fsdp_submodule.module, nn.Linear):
                    num_lora_ab -= 1
                elif isinstance(fsdp_submodule.module, TransformerDecoderLayer):
                    num_transformer_layers -= 1
            assert num_lora_ab == 0
            assert num_transformer_layers == 0
//  [python]
def test_permute_sparse_features(
        self, B: int, T: int, L: int, long_index: bool, has_weight: bool
    ) -> None:
        index_dtype = torch.int64 if long_index else torch.int32
        lengths = torch.randint(low=1, high=L, size=(T, B)).type(index_dtype)
        weights = torch.rand(int(lengths.sum().item())).float() if has_weight else None
        indices = torch.randint(
            low=1,
            high=int(1e5),
            size=cast(Tuple[int, ...], (lengths.sum().item(),)),
        ).type(index_dtype)
        permute_list = list(range(T))
        random.shuffle(permute_list)
        permute = torch.IntTensor(permute_list)

        (
            permuted_lengths_cpu,
            permuted_indices_cpu,
            permuted_weights_cpu,
        ) = torch.ops.fbgemm.permute_sparse_features(permute, lengths, indices, weights)
        (
            permuted_lengths_ref,
            permuted_indices_ref,
            permuted_weights_ref,
            # pyre-fixme[6]: For 4th param expected `LongTensor` but got `Tensor`.
        ) = permute_indices_ref_(lengths, indices, weights, permute.long())
        torch.testing.assert_allclose(permuted_indices_cpu, permuted_indices_ref)
        torch.testing.assert_allclose(permuted_lengths_cpu, permuted_lengths_ref)
        if has_weight:
            torch.testing.assert_allclose(permuted_weights_cpu, permuted_weights_ref)
        else:
            assert permuted_weights_cpu is None and permuted_weights_ref is None

        if gpu_available:
            (
                permuted_lengths_gpu,
                permuted_indices_gpu,
                permuted_weights_gpu,
            ) = torch.ops.fbgemm.permute_sparse_features(
                permute.cuda(),
                lengths.cuda(),
                indices.cuda(),
                weights.cuda() if has_weight and weights is not None else None,
            )
            torch.testing.assert_allclose(permuted_indices_gpu.cpu(), permuted_indices_cpu)
            torch.testing.assert_allclose(permuted_lengths_gpu.cpu(), permuted_lengths_cpu)
            if has_weight:
                torch.testing.assert_allclose(
                    permuted_weights_gpu.cpu(), permuted_weights_cpu
                )
            else:
                assert permuted_weights_gpu is None
//  [python]
def test_permute_sparse_features(
        self, B: int, T: int, L: int, long_index: bool, has_weight: bool
    ) -> None:
        index_dtype = torch.int64 if long_index else torch.int32
        lengths = torch.randint(low=1, high=L, size=(T, B)).type(index_dtype)
        weights = torch.rand(int(lengths.sum().item())).float() if has_weight else None
        indices = torch.randint(
            low=1,
            high=int(1e5),
            size=cast(Tuple[int, ...], (lengths.sum().item(),)),
        ).type(index_dtype)
        permute_list = list(range(T))
        random.shuffle(permute_list)
        permute = torch.IntTensor(permute_list)

        (
            permuted_lengths_cpu,
            permuted_indices_cpu,
            permuted_weights_cpu,
        ) = torch.ops.fbgemm.permute_sparse_features(permute, lengths, indices, weights)
        (
            permuted_lengths_ref,
            permuted_indices_ref,
            permuted_weights_ref,
            # pyre-fixme[6]: For 4th param expected `LongTensor` but got `Tensor`.
        ) = permute_indices_ref_(lengths, indices, weights, permute.long())
        torch.testing.assert_close(permuted_indices_cpu, permuted_indices_ref)
        torch.testing.assert_close(permuted_lengths_cpu, permuted_lengths_ref)
        if has_weight:
            torch.testing.assert_close(permuted_weights_cpu, permuted_weights_ref)
        else:
            assert permuted_weights_cpu is None and permuted_weights_ref is None

        if gpu_available:
            (
                permuted_lengths_gpu,
                permuted_indices_gpu,
                permuted_weights_gpu,
            ) = torch.ops.fbgemm.permute_sparse_features(
                permute.cuda(),
                lengths.cuda(),
                indices.cuda(),
                weights.cuda() if has_weight and weights is not None else None,
            )
            torch.testing.assert_close(permuted_indices_gpu.cpu(), permuted_indices_cpu)
            torch.testing.assert_close(permuted_lengths_gpu.cpu(), permuted_lengths_cpu)
            if has_weight:
                torch.testing.assert_close(
                    permuted_weights_gpu.cpu(), permuted_weights_cpu
                )
            else:
                assert permuted_weights_gpu is None
//  [python]
def _test_multiclass_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: int,
        compute_result: Optional[torch.Tensor] = None,
    ) -> None:

        device = "cpu"
        if torch.cuda.is_available():
            device = "cuda"

        # get sklearn compute result if none given
        if compute_result is None:
            compute_result = self._get_sklearn_equivalent(input, target, device)

        # Get torcheval compute result
        te_compute_result = multiclass_auprc(
            input.to(device=device),
            target.to(device=device),
            num_classes=num_classes,
            average=None,
        )

        # test no average
        torch.testing.assert_allclose(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_multiclass_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: int,
        compute_result: Optional[torch.Tensor] = None,
    ) -> None:

        device = "cpu"
        if torch.cuda.is_available():
            device = "cuda"

        # get sklearn compute result if none given
        if compute_result is None:
            compute_result = self._get_sklearn_equivalent(input, target, device)

        # Get torcheval compute result
        te_compute_result = multiclass_auprc(
            input.to(device=device),
            target.to(device=device),
            num_classes=num_classes,
            average=None,
        )

        # test no average
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_permute_sparse_features_with_repeats(
        self, B: int, T: int, L: int, long_index: bool, has_weight: bool
    ) -> None:
        index_dtype = torch.int64 if long_index else torch.int32
        lengths = torch.randint(low=1, high=L, size=(T, B)).type(index_dtype)
        weights = torch.rand(int(lengths.sum().item())).float() if has_weight else None
        indices = torch.randint(
            low=1,
            high=int(1e5),
            size=cast(Tuple[int, ...], (lengths.sum().item(),)),
        ).type(index_dtype)
        permute_list = list(range(T))

        num_repeats = random.randint(0, T)
        for _ in range(num_repeats):
            permute_list.append(random.randint(0, T - 1))

        random.shuffle(permute_list)
        permute = torch.IntTensor(permute_list)

        (
            permuted_lengths_cpu,
            permuted_indices_cpu,
            permuted_weights_cpu,
        ) = torch.ops.fbgemm.permute_sparse_features(permute, lengths, indices, weights)
        (
            permuted_lengths_ref,
            permuted_indices_ref,
            permuted_weights_ref,
            # pyre-fixme[6]: For 4th param expected `LongTensor` but got `Tensor`.
        ) = permute_indices_ref_(lengths, indices, weights, permute.long())
        t.assert_close(permuted_indices_cpu, permuted_indices_ref)
        t.assert_close(permuted_lengths_cpu, permuted_lengths_ref)
        if has_weight:
            t.assert_close(permuted_weights_cpu, permuted_weights_ref)
        else:
            assert permuted_weights_cpu is None and permuted_weights_ref is None

        if gpu_available:
            (
                permuted_lengths_gpu,
                permuted_indices_gpu,
                permuted_weights_gpu,
            ) = torch.ops.fbgemm.permute_sparse_features(
                permute.cuda(),
                lengths.cuda(),
                indices.cuda(),
                weights.cuda() if has_weight and weights is not None else None,
            )
            t.assert_close(permuted_indices_gpu.cpu(), permuted_indices_cpu)
            t.assert_close(permuted_lengths_gpu.cpu(), permuted_lengths_cpu)
            if has_weight:
                t.assert_close(
                    permuted_weights_gpu.cpu(), permuted_weights_cpu
                )
            else:
                assert permuted_weights_cpu is None
//  [python]
def test_permute_sparse_features_with_repeats(
        self, B: int, T: int, L: int, long_index: bool, has_weight: bool
    ) -> None:
        index_dtype = torch.int64 if long_index else torch.int32
        lengths = torch.randint(low=1, high=L, size=(T, B)).type(index_dtype)
        weights = torch.rand(int(lengths.sum().item())).float() if has_weight else None
        indices = torch.randint(
            low=1,
            high=int(1e5),
            size=cast(Tuple[int, ...], (lengths.sum().item(),)),
        ).type(index_dtype)
        permute_list = list(range(T))

        num_repeats = random.randint(0, T)
        for _ in range(num_repeats):
            permute_list.append(random.randint(0, T - 1))

        random.shuffle(permute_list)
        permute = torch.IntTensor(permute_list)

        (
            permuted_lengths_cpu,
            permuted_indices_cpu,
            permuted_weights_cpu,
        ) = torch.ops.fbgemm.permute_sparse_features(permute, lengths, indices, weights)
        (
            permuted_lengths_ref,
            permuted_indices_ref,
            permuted_weights_ref,
            # pyre-fixme[6]: For 4th param expected `LongTensor` but got `Tensor`.
        ) = permute_indices_ref_(lengths, indices, weights, permute.long())
        torch.testing.assert_close(permuted_indices_cpu, permuted_indices_ref)
        torch.testing.assert_close(permuted_lengths_cpu, permuted_lengths_ref)
        if has_weight:
            torch.testing.assert_close(permuted_weights_cpu, permuted_weights_ref)
        else:
            assert permuted_weights_cpu is None and permuted_weights_ref is None

        if gpu_available:
            (
                permuted_lengths_gpu,
                permuted_indices_gpu,
                permuted_weights_gpu,
            ) = torch.ops.fbgemm.permute_sparse_features(
                permute.cuda(),
                lengths.cuda(),
                indices.cuda(),
                weights.cuda() if has_weight and weights is not None else None,
            )
            torch.testing.assert_close(permuted_indices_gpu.cpu(), permuted_indices_cpu)
            torch.testing.assert_close(permuted_lengths_gpu.cpu(), permuted_lengths_cpu)
            if has_weight:
                torch.testing.assert_close(
                    permuted_weights_gpu.cpu(), permuted_weights_cpu
                )
            else:
                assert permuted_weights_cpu is None
//  [python]
def test_binary_auprc_docstring_examples(self) -> None:
        """This test also tests whether the shape is preserved for 1D and 2D examples with ntasks=1"""
        input = torch.tensor([0.1, 0.5, 0.7, 0.8])
        target = torch.tensor([1, 0, 1, 1])
        compute_result = torch.tensor(0.9167)

        t.assert_close(
            binary_auprc(input, target),
            compute_result,
            atol=1e-4,
            rtol=1e-4,
        )

        input = torch.tensor([[1, 1, 1, 0]])
        target = torch.tensor([[1, 0, 1, 0]])
        compute_result = torch.tensor([0.6667])

        t.assert_close(
            binary_auprc(input, target),
            compute_result,
            atol=1e-4,
            rtol=1e-4,
        )

        input = torch.tensor([[0.1, 0.5, 0.7, 0.8], [1, 1, 1, 0]])
        target = torch.tensor([[1, 0, 1, 1], [1, 0, 1, 0]])
        compute_result = torch.tensor([0.9167, 0.6667])

        t.assert_close(
            binary_auprc(input, target, num_tasks=2),
            compute_result,
            atol=1e-4,
            rtol=1e-4,
        )