jiahuimbzuai/CoQuIR huggingface dataset:
//  [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)

        torch.testing.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])

        torch.testing.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])

        torch.testing.assert_close(
            binary_auprc(input, target, num_tasks=2),
            compute_result,
            atol=1e-4,
            rtol=1e-4,
        )
//  [python]
def test_multiclass_auprc_options(self) -> None:
        # average = macro, num_classes not given
        input, target = self._get_rand_inputs_multiclass(5, BATCH_SIZE)
        compute_result = torch.mean(self._get_sklearn_equivalent(input, target))
        te_compute_result = multiclass_auprc(input, target, average="macro")
        t.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # average = macro (not given), num_classes given
        input, target = self._get_rand_inputs_multiclass(5, BATCH_SIZE)
        compute_result = torch.mean(self._get_sklearn_equivalent(input, target))
        te_compute_result = multiclass_auprc(input, target)
        t.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # average = none
        input, target = self._get_rand_inputs_multiclass(5, BATCH_SIZE)
        compute_result = self._get_sklearn_equivalent(input, target)
        te_compute_result = multiclass_auprc(input, target, average="none")
        t.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
        te_compute_result = multiclass_auprc(input, target, average=None)
        t.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_multiclass_auprc_options(self) -> None:
        # average = macro, num_classes not given
        input, target = self._get_rand_inputs_multiclass(5, BATCH_SIZE)
        compute_result = torch.mean(self._get_sklearn_equivalent(input, target))
        te_compute_result = multiclass_auprc(input, target, average="macro")
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # average = macro (not given), num_classes given
        input, target = self._get_rand_inputs_multiclass(5, BATCH_SIZE)
        compute_result = torch.mean(self._get_sklearn_equivalent(input, target))
        te_compute_result = multiclass_auprc(input, target)
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # average = none
        input, target = self._get_rand_inputs_multiclass(5, BATCH_SIZE)
        compute_result = self._get_sklearn_equivalent(input, target)
        te_compute_result = multiclass_auprc(input, target, average="none")
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
        te_compute_result = multiclass_auprc(input, target, average=None)
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_multiclass_auprc_docstring_examples(self) -> None:
        input = torch.tensor([[0.5647, 0.2726], [0.9143, 0.1895], [0.7782, 0.3082]])
        target = torch.tensor([0, 1, 0])
        output = torch.tensor([0.5833, 0.3333])
        result = multiclass_auprc(input, target, average=None)
        torch.testing.assert_allclose(
            result,
            output,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )

        avg_result = multiclass_auprc(input, target)
        avg_output = torch.tensor(0.4583)
        torch.testing.assert_allclose(
            avg_result,
            avg_output,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )

        input = torch.tensor([[0.1, 1], [0.5, 1], [0.7, 1], [0.8, 0]])
        target = torch.tensor([1, 0, 0, 1])
        result = multiclass_auprc(input, target, 2, average=None)
        output = torch.tensor([0.5833, 0.4167])
        torch.testing.assert_allclose(
            avg_result,
            avg_output,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )
//  [python]
def test_multiclass_auprc_docstring_examples(self) -> None:
        input = torch.tensor([[0.5647, 0.2726], [0.9143, 0.1895], [0.7782, 0.3082]])
        target = torch.tensor([0, 1, 0])
        output = torch.tensor([0.5833, 0.3333])
        result = multiclass_auprc(input, target, average=None)
        torch.testing.assert_close(
            result,
            output,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )

        avg_result = multiclass_auprc(input, target)
        avg_output = torch.tensor(0.4583)
        torch.testing.assert_close(
            avg_result,
            avg_output,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )

        input = torch.tensor([[0.1, 1], [0.5, 1], [0.7, 1], [0.8, 0]])
        target = torch.tensor([1, 0, 0, 1])
        result = multiclass_auprc(input, target, 2, average=None)
        output = torch.tensor([0.5833, 0.4167])
        torch.testing.assert_close(
            avg_result,
            avg_output,
            equal_nan=True,
            atol=1e-4,
            rtol=1e-3,
        )
//  [python]
def _test_multilabel_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_labels: int,
        average: Optional[str] = "macro",
        compute_result: Optional[torch.Tensor] = None,
    ) -> None:

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

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

        # Get torcheval compute result
        te_compute_result = multilabel_auprc(
            input.to(device=device),
            target.to(device=device),
            num_labels=num_labels,
            average=average,
        )

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

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

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

        # Get torcheval compute result
        te_compute_result = multilabel_auprc(
            input.to(device=device),
            target.to(device=device),
            num_labels=num_labels,
            average=average,
        )

        # test no average
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_multilabel_auprc_unspecified_average(self, num_labels: int) -> None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        input, target = self._get_rand_inputs_multilabel(num_labels, BATCH_SIZE)

        compute_result = self._get_sklearn_equivalent(input, target, device=device)

        te_compute_result = multilabel_auprc(
            input.to(device=device),
            target.to(device=device),
            num_labels=num_labels,
        )
        torch.testing.assert_allclose(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_multilabel_auprc_unspecified_average(self, num_labels: int) -> None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        input, target = self._get_rand_inputs_multilabel(num_labels, BATCH_SIZE)

        compute_result = self._get_sklearn_equivalent(input, target, device=device)

        te_compute_result = multilabel_auprc(
            input.to(device=device),
            target.to(device=device),
            num_labels=num_labels,
        )
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_space_to_depth(n, b, h, w):
    tensor_input = torch.rand((n, b, h, w))
    out1 = focus_transform(tensor_input)
    out2 = space_to_depth(tensor_input)
    torch.testing.assert_allclose(out2, out1)
//  [python]
def test_space_to_depth(n, b, h, w):
    tensor_input = torch.rand((n, b, h, w))
    out1 = focus_transform(tensor_input)
    out2 = space_to_depth(tensor_input)
    torch.testing.assert_close(out2, out1)
//  [python]
def _totensor(self, test_scripting):
        padding_value = 0
        transform = transforms.ToTensor(padding_value=padding_value)
        if test_scripting:
            transform = torch.jit.script(transform)
        input = [[1, 2], [1, 2, 3]]

        actual = transform(input)
        expected = torch.tensor([[1, 2, 0], [1, 2, 3]], dtype=torch.long)
        torch.testing.assert_allclose(actual, expected)

        input = [1, 2]
        actual = transform(input)
        expected = torch.tensor([1, 2], dtype=torch.long)
        torch.testing.assert_allclose(actual, expected)
//  [python]
def _totensor(self, test_scripting):
        padding_value = 0
        transform = transforms.ToTensor(padding_value=padding_value)
        if test_scripting:
            transform = torch.jit.script(transform)
        input = [[1, 2], [1, 2, 3]]

        actual = transform(input)
        expected = torch.tensor([[1, 2, 0], [1, 2, 3]], dtype=torch.long)
        torch.testing.assert_close(actual, expected)

        input = [1, 2]
        actual = transform(input)
        expected = torch.tensor([1, 2], dtype=torch.long)
        torch.testing.assert_close(actual, expected)
//  [python]
def _sequential(self, test_scripting):
        max_seq_len = 3
        padding_val = 0
        transform = transforms.Sequential(
            transforms.Truncate(max_seq_len=max_seq_len),
            transforms.ToTensor(padding_value=padding_val, dtype=torch.long),
        )

        if test_scripting:
            transform = torch.jit.script(transform)

        input = [[1, 2, 3], [1, 2, 3]]

        actual = transform(input)
        expected = torch.tensor(input)
        torch.testing.assert_allclose(actual, expected)
//  [python]
def _sequential(self, test_scripting):
        max_seq_len = 3
        padding_val = 0
        transform = transforms.Sequential(
            transforms.Truncate(max_seq_len=max_seq_len),
            transforms.ToTensor(padding_value=padding_val, dtype=torch.long),
        )

        if test_scripting:
            transform = torch.jit.script(transform)

        input = [[1, 2, 3], [1, 2, 3]]

        actual = transform(input)
        expected = torch.tensor(input)
        torch.testing.assert_close(actual, expected)
//  [python]
def ort_validate(self, onnx_io, inputs, outputs):

        inputs, _ = torch.jit._flatten(inputs)
        outputs, _ = torch.jit._flatten(outputs)

        def to_numpy(tensor):
            if tensor.requires_grad:
                return tensor.detach().cpu().numpy()
            else:
                return tensor.cpu().numpy()

        inputs = list(map(to_numpy, inputs))
        outputs = list(map(to_numpy, outputs))

        ort_session = onnxruntime.InferenceSession(onnx_io.getvalue(), providers=onnxruntime.get_available_providers())
        # compute onnxruntime output prediction
        ort_inputs = {ort_session.get_inputs()[i].name: inpt for i, inpt in enumerate(inputs)}
        ort_outs = ort_session.run(None, ort_inputs)

        for i in range(0, len(outputs)):
            torch.allclose(outputs[i], ort_outs[i], rtol=1e-03, atol=1e-05)
//  [python]
def ort_validate(self, onnx_io, inputs, outputs):

        inputs, _ = torch.jit._flatten(inputs)
        outputs, _ = torch.jit._flatten(outputs)

        def to_numpy(tensor):
            if tensor.requires_grad:
                return tensor.detach().cpu().numpy()
            else:
                return tensor.cpu().numpy()

        inputs = list(map(to_numpy, inputs))
        outputs = list(map(to_numpy, outputs))

        ort_session = onnxruntime.InferenceSession(onnx_io.getvalue(), providers=onnxruntime.get_available_providers())
        # compute onnxruntime output prediction
        ort_inputs = {ort_session.get_inputs()[i].name: inpt for i, inpt in enumerate(inputs)}
        ort_outs = ort_session.run(None, ort_inputs)

        for i in range(0, len(outputs)):
            torch.testing.assert_close(outputs[i], ort_outs[i], rtol=1e-03, atol=1e-05)
//  [python]
def test_multilabel_auprc_options(self) -> None:
        # average = macro, num_labels not given
        input, target = self._get_rand_inputs_multilabel(5, BATCH_SIZE)
        compute_result = self._get_sklearn_equivalent(input, target)
        te_compute_result = multilabel_auprc(input, target, average="macro")
        self.assertEqual(te_compute_result.shape, compute_result.shape)
        self.assertTrue(torch.allclose(te_compute_result, compute_result, atol=1e-8, rtol=1e-5))

        # average = macro (not given), num_labels given
        input, target = self._get_rand_inputs_multilabel(5, BATCH_SIZE)
        compute_result = self._get_sklearn_equivalent(input, target)
        te_compute_result = multilabel_auprc(input, target, num_labels=5)
        self.assertEqual(te_compute_result.shape, compute_result.shape)
        self.assertTrue(torch.allclose(te_compute_result, compute_result, atol=1e-8, rtol=1e-5))

        # average = macro (not given), num_labels (not given)
        input, target = self._get_rand_inputs_multilabel(5, BATCH_SIZE)
        compute_result = self._get_sklearn_equivalent(input, target)
        te_compute_result = multilabel_auprc(input, target)
        self.assertEqual(te_compute_result.shape, compute_result.shape)
        self.assertTrue(torch.allclose(te_compute_result, compute_result, atol=1e-8, rtol=1e-5))

        # average = none, num_labels not given
        input, target = self._get_rand_inputs_multilabel(5, BATCH_SIZE)
        compute_result = self._get_sklearn_equivalent(input, target, average="none")
        te_compute_result = multilabel_auprc(input, target, average="none")
        self.assertEqual(te_compute_result.shape, compute_result.shape)
        self.assertTrue(torch.allclose(te_compute_result, compute_result, atol=1e-8, rtol=1e-5))
        te_compute_result = multilabel_auprc(input, target, average=None)
        self.assertEqual(te_compute_result.shape, compute_result.shape)
        self.assertTrue(torch.allclose(te_compute_result, compute_result, atol=1e-8, rtol=1e-5))
//  [python]
def test_multilabel_auprc_options(self) -> None:
        # average = macro, num_labels not given
        input, target = self._get_rand_inputs_multilabel(5, BATCH_SIZE)
        compute_result = self._get_sklearn_equivalent(input, target)
        te_compute_result = multilabel_auprc(input, target, average="macro")
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # average = macro (not given), num_labels given
        input, target = self._get_rand_inputs_multilabel(5, BATCH_SIZE)
        compute_result = self._get_sklearn_equivalent(input, target)
        te_compute_result = multilabel_auprc(input, target, num_labels=5)
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # average = macro (not given), num_labels (not given)
        input, target = self._get_rand_inputs_multilabel(5, BATCH_SIZE)
        compute_result = self._get_sklearn_equivalent(input, target)
        te_compute_result = multilabel_auprc(input, target)
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # average = none, num_labels not given
        input, target = self._get_rand_inputs_multilabel(5, BATCH_SIZE)
        compute_result = self._get_sklearn_equivalent(input, target, average="none")
        te_compute_result = multilabel_auprc(input, target, average="none")
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
        te_compute_result = multilabel_auprc(input, target, average=None)
        torch.testing.assert_close(
            te_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def torch_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    nvfuser_result = fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))
    torch_result = nvf_op.reference(*sample.args, **sample.kwargs)

    if isinstance(nvfuser_result, Exception):
        raise nvfuser_result

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]

    # TODO If dtype is fp16 or bf16, skip dtype check because nvfuser promotes to fp32 but does not return original dtype.
    # TODO Add specific dtype tolerances
    t.assert_close(
        nvfuser_result, torch_result, equal_nan=True, atol=1e-3, rtol=0
    )
//  [python]
def torch_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    nvfuser_result = fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))
    torch_result = nvf_op.reference(*sample.args, **sample.kwargs)

    if isinstance(nvfuser_result, Exception):
        raise nvfuser_result

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]

    # TODO If dtype is fp16 or bf16, skip dtype check because nvfuser promotes to fp32 but does not return original dtype.
    # TODO Add specific dtype tolerances
    torch.testing.assert_close(
        nvfuser_result, torch_result, equal_nan=True, atol=1e-3, rtol=0
    )
//  [python]
def __test_on_device(self, device):
        layer = LeNet5(5).to(device)
        x = torch.randn(2, 1, 32, 32, device=device)
        saved_file = os.path.join(self.temp_dir, "test_inputs.pth")
        torch.save([x, layer.state_dict()], saved_file)

        with module_numberic_checker(layer, mode="save"):
            out = layer(x)

        # test load
        layer2 = LeNet5(5).to(device)
        loaded_x, layer_state_dict = torch.load(saved_file)
        layer_state_dict = move_to_device(layer_state_dict, device)
        loaded_x = move_to_device(loaded_x, device)
        with module_numberic_checker(layer2, mode="load") as no_allclose:
            self.assertEqual(len(no_allclose), 0)
            layer2(loaded_x)
        self.assertNotEqual(len(no_allclose), 0)
        for func_name, err_msg in no_allclose.items():
            print(func_name, err_msg)
        layer2.load_state_dict(layer_state_dict)

        with module_numberic_checker(layer2, mode="load") as no_allclose:
            self.assertEqual(len(no_allclose), 0)
            out2 = layer2(loaded_x)
        self.assertEqual(len(no_allclose), 0)
        torch.testing.assert_allclose(out2, out)
//  [python]
def __test_on_device(self, device):
        layer = LeNet5(5).to(device)
        x = torch.randn(2, 1, 32, 32, device=device)
        saved_file = os.path.join(self.temp_dir, "test_inputs.pth")
        torch.save([x, layer.state_dict()], saved_file)

        with module_numberic_checker(layer, mode="save"):
            out = layer(x)

        # test load
        layer2 = LeNet5(5).to(device)
        loaded_x, layer_state_dict = torch.load(saved_file)
        layer_state_dict = move_to_device(layer_state_dict, device)
        loaded_x = move_to_device(loaded_x, device)
        with module_numberic_checker(layer2, mode="load") as no_allclose:
            self.assertEqual(len(no_allclose), 0)
            layer2(loaded_x)
        self.assertNotEqual(len(no_allclose), 0)
        for func_name, err_msg in no_allclose.items():
            print(func_name, err_msg)
        layer2.load_state_dict(layer_state_dict)

        with module_numberic_checker(layer2, mode="load") as no_allclose:
            self.assertEqual(len(no_allclose), 0)
            out2 = layer2(loaded_x)
        self.assertEqual(len(no_allclose), 0)
        torch.testing.assert_close(out2, out)
//  [python]
def test_multilabel_auprc_docstring_example(self) -> None:
        input = torch.tensor(
            [
                [0.75, 0.05, 0.35],
                [0.45, 0.75, 0.05],
                [0.05, 0.55, 0.75],
                [0.05, 0.65, 0.05],
            ]
        )
        target = torch.tensor([[1, 0, 1], [0, 0, 0], [0, 1, 1], [1, 1, 1]])

        compute_result_average_none = torch.tensor([0.7500, 0.5833, 0.9167])
        t.assert_close(
            multilabel_auprc(input, target, num_labels=3, average=None),
            compute_result_average_none,
            atol=1e-4,
            rtol=1e-4,
        )
        t.assert_close(
            multilabel_auprc(input, target, average=None),
            compute_result_average_none,
            atol=1e-4,
            rtol=1e-4,
        )

        compute_result_average_macro = torch.tensor(0.7500)
        t.assert_close(
            multilabel_auprc(input, target, num_labels=3, average="macro"),
            compute_result_average_macro,
            atol=1e-4,
            rtol=1e-4,
        )
        t.assert_close(
            multilabel_auprc(input, target, num_labels=3),
            compute_result_average_macro,
            atol=1e-4,
            rtol=1e-4,
        )
        t.assert_close(
            multilabel_auprc(input, target, average="macro"),
            compute_result_average_macro,
            atol=1e-4,
            rtol=1e-4,
        )
        t.assert_close(
            multilabel_auprc(input, target),
            compute_result_average_macro,
            atol=1e-4,
            rtol=1e-4,
        )

        input = torch.tensor([[0.1, 0, 0], [0, 1, 0], [0.1, 0.2, 0.7], [0, 0, 1]])
        target = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 1]])
        compute_result_average_none = torch.tensor([0.5000, 1.0000, 1.0000])
        t.assert_close(
            multilabel_auprc(input, target, average=None),
            compute_result_average_none,
            atol=1e-4,
            rtol=1e-4,
        )
//  [python]
def test_multilabel_auprc_docstring_example(self) -> None:
        input = torch.tensor(
            [
                [0.75, 0.05, 0.35],
                [0.45, 0.75, 0.05],
                [0.05, 0.55, 0.75],
                [0.05, 0.65, 0.05],
            ]
        )
        target = torch.tensor([[1, 0, 1], [0, 0, 0], [0, 1, 1], [1, 1, 1]])

        compute_result_average_none = torch.tensor([0.7500, 0.5833, 0.9167])
        torch.testing.assert_close(
            multilabel_auprc(input, target, num_labels=3, average=None),
            compute_result_average_none,
            atol=1e-4,
            rtol=1e-4,
        )
        torch.testing.assert_close(
            multilabel_auprc(input, target, average=None),
            compute_result_average_none,
            atol=1e-4,
            rtol=1e-4,
        )

        compute_result_average_macro = torch.tensor(0.7500)
        torch.testing.assert_close(
            multilabel_auprc(input, target, num_labels=3, average="macro"),
            compute_result_average_macro,
            atol=1e-4,
            rtol=1e-4,
        )
        torch.testing.assert_close(
            multilabel_auprc(input, target, num_labels=3),
            compute_result_average_macro,
            atol=1e-4,
            rtol=1e-4,
        )
        torch.testing.assert_close(
            multilabel_auprc(input, target, average="macro"),
            compute_result_average_macro,
            atol=1e-4,
            rtol=1e-4,
        )
        torch.testing.assert_close(
            multilabel_auprc(input, target),
            compute_result_average_macro,
            atol=1e-4,
            rtol=1e-4,
        )

        input = torch.tensor([[0.1, 0, 0], [0, 1, 0], [0.1, 0.2, 0.7], [0, 0, 1]])
        target = torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 1]])
        compute_result_average_none = torch.tensor([0.5000, 1.0000, 1.0000])
        torch.testing.assert_close(
            multilabel_auprc(input, target, average=None),
            compute_result_average_none,
            atol=1e-4,
            rtol=1e-4,
        )
//  [python]
def jax_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    nvfuser_result = fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))

    jax_sample = sample.jax()
    jax_result = nvf_op.reference(*jax_sample.args, **jax_sample.kwargs)

    # NOTE: this strange unpacking is to handle NumPy's and JAX's sometimes odd
    #   number vs. array representation. In particular, NumPy can mimic
    #   Python numbers, but `asarray` doesn't understand this mimicry
    np_array = np.array(jax_result)
    if np_array.shape == ():
        jax_result = torch.tensor(np_array.item(), device="cuda")
    else:
        jax_result = torch.asarray(np_array, device="cuda")

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]

    # NOTE: dtype is not checked because jax will translate int64, float64, and complex128 to int32, float32 and complex64
    t.assert_close(
        nvfuser_result, jax_result, equal_nan=True, atol=1e-3, rtol=0, check_dtype=False
    )
//  [python]
def jax_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args, **sample.kwargs)
    nvfuser_result = fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))

    jax_sample = sample.jax()
    jax_result = nvf_op.reference(*jax_sample.args, **jax_sample.kwargs)

    # NOTE: this strange unpacking is to handle NumPy's and JAX's sometimes odd
    #   number vs. array representation. In particular, NumPy can mimic
    #   Python numbers, but `asarray` doesn't understand this mimicry
    np_array = np.array(jax_result)
    if np_array.shape == ():
        jax_result = torch.tensor(np_array.item(), device="cuda")
    else:
        jax_result = torch.asarray(np_array, device="cuda")

    if len(nvfuser_result) == 1:
        nvfuser_result = nvfuser_result[0]

    # NOTE: dtype is not checked because jax will translate int64, float64, and complex128 to int32, float32 and complex64
    torch.testing.assert_close(
        nvfuser_result, jax_result, equal_nan=True, atol=1e-3, rtol=0, check_dtype=False
    )
//  [python]
def _pad_transform(self, test_scripting):
        """
        Test padding transform on 1D and 2D tensors.
        When max_length < tensor length at dim -1, this should be a no-op.
        Otherwise the tensor should be padded to max_length in dim -1.
        """

        input_1d_tensor = torch.ones(5)
        input_2d_tensor = torch.ones((8, 5))
        pad_long = transforms.PadTransform(max_length=7, pad_value=0)
        if test_scripting:
            pad_long = torch.jit.script(pad_long)
        padded_1d_tensor_actual = pad_long(input_1d_tensor)
        padded_1d_tensor_expected = torch.cat([torch.ones(5), torch.zeros(2)])
        torch.testing.assert_allclose(
            padded_1d_tensor_actual,
            padded_1d_tensor_expected,
            msg=f"actual: {padded_1d_tensor_actual}, expected: {padded_1d_tensor_expected}",
        )

        padded_2d_tensor_actual = pad_long(input_2d_tensor)
        padded_2d_tensor_expected = torch.cat([torch.ones(8, 5), torch.zeros(8, 2)], axis=-1)
        torch.testing.assert_allclose(
            padded_2d_tensor_actual,
            padded_2d_tensor_expected,
            msg=f"actual: {padded_2d_tensor_actual}, expected: {padded_2d_tensor_expected}",
        )

        pad_short = transforms.PadTransform(max_length=3, pad_value=0)
        if test_scripting:
            pad_short = torch.jit.script(pad_short)
        padded_1d_tensor_actual = pad_short(input_1d_tensor)
        padded_1d_tensor_expected = input_1d_tensor
        torch.testing.assert_allclose(
            padded_1d_tensor_actual,
            padded_1d_tensor_expected,
            msg=f"actual: {padded_1d_tensor_actual}, expected: {padded_1d_tensor_expected}",
        )

        padded_2d_tensor_actual = pad_short(input_2d_tensor)
        padded_2d_tensor_expected = input_2d_tensor
        torch.testing.assert_allclose(
            padded_2d_tensor_actual,
            padded_2d_tensor_expected,
            msg=f"actual: {padded_2d_tensor_actual}, expected: {padded_2d_tensor_expected}",
        )
//  [python]
def _pad_transform(self, test_scripting):
        """
        Test padding transform on 1D and 2D tensors.
        When max_length < tensor length at dim -1, this should be a no-op.
        Otherwise the tensor should be padded to max_length in dim -1.
        """

        input_1d_tensor = torch.ones(5)
        input_2d_tensor = torch.ones((8, 5))
        pad_long = transforms.PadTransform(max_length=7, pad_value=0)
        if test_scripting:
            pad_long = torch.jit.script(pad_long)
        padded_1d_tensor_actual = pad_long(input_1d_tensor)
        padded_1d_tensor_expected = torch.cat([torch.ones(5), torch.zeros(2)])
        torch.testing.assert_close(
            padded_1d_tensor_actual,
            padded_1d_tensor_expected,
            msg=f"actual: {padded_1d_tensor_actual}, expected: {padded_1d_tensor_expected}",
        )

        padded_2d_tensor_actual = pad_long(input_2d_tensor)
        padded_2d_tensor_expected = torch.cat([torch.ones(8, 5), torch.zeros(8, 2)], axis=-1)
        torch.testing.assert_close(
            padded_2d_tensor_actual,
            padded_2d_tensor_expected,
            msg=f"actual: {padded_2d_tensor_actual}, expected: {padded_2d_tensor_expected}",
        )

        pad_short = transforms.PadTransform(max_length=3, pad_value=0)
        if test_scripting:
            pad_short = torch.jit.script(pad_short)
        padded_1d_tensor_actual = pad_short(input_1d_tensor)
        padded_1d_tensor_expected = input_1d_tensor
        torch.testing.assert_close(
            padded_1d_tensor_actual,
            padded_1d_tensor_expected,
            msg=f"actual: {padded_1d_tensor_actual}, expected: {padded_1d_tensor_expected}",
        )

        padded_2d_tensor_actual = pad_short(input_2d_tensor)
        padded_2d_tensor_expected = input_2d_tensor
        torch.testing.assert_close(
            padded_2d_tensor_actual,
            padded_2d_tensor_expected,
            msg=f"actual: {padded_2d_tensor_actual}, expected: {padded_2d_tensor_expected}",
        )
//  [python]
def python_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    # python reference function does not accept keyword arguments
    assert len(sample.kwargs) == 0

    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args)
    nvfuser_result = fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))

    # expect only single result from function
    assert len(nvfuser_result) == 1

    # convert tensor arguments into flat, python lists
    python_sample = sample.python()

    # apply reference to python lists
    python_result = map(nvf_op.reference, *python_sample.args)

    # create pytorch tensor
    np_array = np.array(list(python_result))
    if np_array.shape == ():
        python_result = torch.tensor(
            np_array.item(), dtype=nvfuser_result[0].dtype, device="cuda"
        )
    else:
        python_result = torch.asarray(
            np_array, dtype=nvfuser_result[0].dtype, device="cuda"
        )

    # reshape flat output tensor into expected shape
    torch.testing.assert_allclose(
        nvfuser_result[0],
        python_result.reshape(nvfuser_result[0].shape),
        equal_nan=True,
        atol=1e-3,
        rtol=0,
    )
//  [python]
def python_correctness_test_fn(fd_fn: Callable, nvf_op: OpInfo, sample: SampleInput):
    # python reference function does not accept keyword arguments
    assert len(sample.kwargs) == 0

    with FusionDefinition() as fd:
        fd_fn(fd, nvf_op, *sample.args)
    nvfuser_result = fd.execute(parse_args_fusion_execution(nvf_op, *sample.args))

    # expect only single result from function
    assert len(nvfuser_result) == 1

    # convert tensor arguments into flat, python lists
    python_sample = sample.python()

    # apply reference to python lists
    python_result = map(nvf_op.reference, *python_sample.args)

    # create pytorch tensor
    np_array = np.array(list(python_result))
    if np_array.shape == ():
        python_result = torch.tensor(
            np_array.item(), dtype=nvfuser_result[0].dtype, device="cuda"
        )
    else:
        python_result = torch.asarray(
            np_array, dtype=nvfuser_result[0].dtype, device="cuda"
        )

    # reshape flat output tensor into expected shape
    torch.testing.assert_close(
        nvfuser_result[0],
        python_result.reshape(nvfuser_result[0].shape),
        equal_nan=True,
        atol=1e-3,
        rtol=0,
    )
//  [python]
def dequantize_test_helper(self, weight_type, quant_type):
        assert quant_type == torch.int8 or quant_type == torch.quint4x2

        lower_bound = -128 if quant_type == torch.int8 else -8
        upper_bound = 127 if quant_type == torch.int8 else 7

        m, n, k = 64, 128, 64
        weights = torch.randint(lower_bound,
                                upper_bound, [k, n],
                                dtype=torch.int8,
                                device='cpu')

        packed_weight = self.pack_int4s(
            weights) if quant_type == torch.quint4x2 else weights
        cuda_weights = self.preprocess_weights_for_mixed_gemm(
            packed_weight, quant_type).to('cuda')
        weights = weights.to('cuda')

        act = torch.eye(m, dtype=weight_type, device='cuda')
        scales = torch.ones([n], dtype=weight_type, device='cuda')

        actual = self.fused_gemm_dq(act, cuda_weights, scales)
        self.assertEqual(actual, weights)
//  [python]
def dequantize_test_helper(self, weight_type, quant_type):
        assert quant_type == torch.int8 or quant_type == torch.quint4x2

        lower_bound = -128 if quant_type == torch.int8 else -8
        upper_bound = 127 if quant_type == torch.int8 else 7

        m, n, k = 64, 128, 64
        weights = torch.randint(lower_bound,
                                upper_bound, [k, n],
                                dtype=torch.int8,
                                device='cpu')

        packed_weight = self.pack_int4s(
            weights) if quant_type == torch.quint4x2 else weights
        cuda_weights = self.preprocess_weights_for_mixed_gemm(
            packed_weight, quant_type).to('cuda')
        weights = weights.to('cuda')

        act = torch.eye(m, dtype=weight_type, device='cuda')
        scales = torch.ones([n], dtype=weight_type, device='cuda')

        actual = self.fused_gemm_dq(act, cuda_weights, scales)
        torch.testing.assert_close(actual,
                                   weights,
                                   atol=0,
                                   rtol=0,
                                   check_dtype=False)
//  [python]
def test_2d_linear(self):
        x = torch.randn(2, 3, 4)
        linear = Linear2D(4, 5)
        old_linear = nn.Linear(4, 5)
        old_linear.load_state_dict(linear.state_dict())

        y = linear(x)
        self.assertEqual(y.size(), (2, 3, 5))

        y1 = old_linear(x)
        self.assertEqual(y1.size(), (2, 3, 5))
        torch.testing.assert_allclose(y, y1)
        Y = torch.randn_like(y)
        loss = torch.nn.functional.mse_loss(y, Y)
        loss1 = torch.nn.functional.mse_loss(y1, Y)
        loss.backward()
        loss1.backward()
        name2grad = {}
        for name, p in linear.named_parameters():
            name2grad[name] = p.grad
        for name, p in old_linear.named_parameters():
            torch.testing.assert_allclose(name2grad[name], p.grad, rtol=1e-4, atol=1e-4)
//  [python]
def test_2d_linear(self):
        x = torch.randn(2, 3, 4)
        linear = Linear2D(4, 5)
        old_linear = nn.Linear(4, 5)
        old_linear.load_state_dict(linear.state_dict())

        y = linear(x)
        self.assertEqual(y.size(), (2, 3, 5))

        y1 = old_linear(x)
        self.assertEqual(y1.size(), (2, 3, 5))
        torch.testing.assert_close(y, y1)
        Y = torch.randn_like(y)
        loss = torch.nn.functional.mse_loss(y, Y)
        loss1 = torch.nn.functional.mse_loss(y1, Y)
        loss.backward()
        loss1.backward()
        name2grad = {}
        for name, p in linear.named_parameters():
            name2grad[name] = p.grad
        for name, p in old_linear.named_parameters():
            torch.testing.assert_close(name2grad[name], p.grad, rtol=1e-4, atol=1e-4)
//  [python]
def test_amp_module(self):
        device = torch.cuda.current_device()
        batch, seq_length = 10, 500
        input_dim, output_dim = 128, 666
        linear0 = torch.nn.Linear(input_dim, output_dim)
        x = torch.randn(batch, seq_length, input_dim, device=device)
        mod = AmpFusedDense(input_dim, output_dim)
        mod.to(device)
        linear0.to(device)
        linear0.zero_grad()
        mod.weight.data.copy_(linear0.weight)
        mod.bias.data.copy_(linear0.bias)
        out = mod(x.reshape(batch * seq_length, -1))
        out = out.reshape(batch, seq_length, -1)
        out1 = linear0(x)
        torch.testing.assert_allclose(out, out1)
        Y = torch.randn_like(out)
        loss = torch.sum(torch.pow(Y - out, 2))
        loss.backward()
        loss = torch.sum(torch.pow(Y - out1, 2))
        loss.backward()
        for (name, p0), p1 in zip(mod.named_parameters(), linear0.parameters()):
            if p0.grad is not None:
                with self.subTest("grad check", name=name):
                    torch.testing.assert_allclose(p0.grad, p1.grad)
//  [python]
def test_amp_module(self):
        device = torch.cuda.current_device()
        batch, seq_length = 10, 500
        input_dim, output_dim = 128, 666
        linear0 = torch.nn.Linear(input_dim, output_dim)
        x = torch.randn(batch, seq_length, input_dim, device=device)
        mod = AmpFusedDense(input_dim, output_dim)
        mod.to(device)
        linear0.to(device)
        linear0.zero_grad()
        mod.weight.data.copy_(linear0.weight)
        mod.bias.data.copy_(linear0.bias)
        out = mod(x.reshape(batch * seq_length, -1))
        out = out.reshape(batch, seq_length, -1)
        out1 = linear0(x)
        torch.testing.assert_close(out, out1)
        Y = torch.randn_like(out)
        loss = torch.sum(torch.pow(Y - out, 2))
        loss.backward()
        loss = torch.sum(torch.pow(Y - out1, 2))
        loss.backward()
        for (name, p0), p1 in zip(mod.named_parameters(), linear0.parameters()):
            if p0.grad is not None:
                with self.subTest("grad check", name=name):
                    torch.testing.assert_allclose(p0.grad, p1.grad)
//  [python]
def test_two_layer(self):
        device = torch.device("cuda")
        dtype = torch.float32
        batch, head_dim, head_num, hidden = 128, 64, 8, 768
        x = torch.randn(batch, head_dim, head_num, hidden, device=device, dtype=dtype)
        layer = MyModule(True, hidden, hidden).to(device)
        layer.reset_parameters()
        layer_fused = MyModule(False, hidden, hidden).to(device)
        layer_fused.load_state_dict(layer.state_dict())
        with torch.cuda.amp.autocast(dtype=dtype):
            y = layer(x)
            y1 = layer_fused(x)
            torch.testing.assert_allclose(y, y1, rtol=1e-4, atol=1e-4)
            Y = torch.randn_like(y)
            loss = torch.nn.functional.mse_loss(y, Y)
            loss1 = torch.nn.functional.mse_loss(y1, Y)
            loss.backward()
            loss1.backward()
            name2grad = {}
            for name, p in layer.named_parameters():
                name2grad[name] = p.grad
            for name, p in layer_fused.named_parameters():
                torch.testing.assert_allclose(name2grad[name], p.grad, rtol=1e-4, atol=1e-4)
//  [python]
def test_two_layer(self):
        device = torch.device("cuda")
        dtype = torch.float32
        batch, head_dim, head_num, hidden = 128, 64, 8, 768
        x = torch.randn(batch, head_dim, head_num, hidden, device=device, dtype=dtype)
        layer = MyModule(True, hidden, hidden).to(device)
        layer.reset_parameters()
        layer_fused = MyModule(False, hidden, hidden).to(device)
        layer_fused.load_state_dict(layer.state_dict())
        with torch.cuda.amp.autocast(dtype=dtype):
            y = layer(x)
            y1 = layer_fused(x)
            torch.testing.assert_close(y, y1, rtol=1e-4, atol=1e-4)
            Y = torch.randn_like(y)
            loss = torch.nn.functional.mse_loss(y, Y)
            loss1 = torch.nn.functional.mse_loss(y1, Y)
            loss.backward()
            loss1.backward()
            name2grad = {}
            for name, p in layer.named_parameters():
                name2grad[name] = p.grad
            for name, p in layer_fused.named_parameters():
                torch.testing.assert_close(name2grad[name], p.grad, rtol=1e-4, atol=1e-4)
//  [python]
def verify(lazy_results, cpu_results, tol=1e-5):
    """Verify the series of losses."""
    logger.debug("lazy_losses = %s", lazy_results)
    logger.debug("cpu_losses = %s", cpu_results)
    for lazy, cpu in zip(lazy_results, cpu_results):
        t.assert_allclose(lazy, cpu, atol=tol, rtol=tol)
//  [python]
def verify(lazy_results, cpu_results, tol=1e-5):
    """Verify the series of losses."""
    logger.debug("lazy_losses = %s", lazy_results)
    logger.debug("cpu_losses = %s", cpu_results)
    for lazy, cpu in zip(lazy_results, cpu_results):
        torch.testing.assert_close(lazy, cpu, atol=tol, rtol=tol)
//  [python]
def check_hf_model_equivalence(model1: PreTrainedModel,
                               model2: PreTrainedModel):
    # Checks that two huggingface models are equivalent (config and
    # parameters)
    expected_model_config_dict = model1.config.to_dict()
    new_model_config_dict = model2.config.to_dict()

    # this key just says the folder it was loaded from, which is a tmp dir during pytest
    del expected_model_config_dict['_name_or_path']
    del new_model_config_dict['_name_or_path']

    assert expected_model_config_dict == new_model_config_dict
    assert sum(p.numel() for p in model1.parameters()) == sum(
        p.numel() for p in model2.parameters())
    assert all(
        type(module1) == type(module2)
        for module1, module2 in zip(model1.modules(), model2.modules()))

    for p1, p2 in zip(model1.parameters(), model2.parameters()):
        torch.testing.assert_allclose(p1, p2)
//  [python]
def check_hf_model_equivalence(model1: PreTrainedModel,
                               model2: PreTrainedModel):
    # Checks that two huggingface models are equivalent (config and
    # parameters)
    expected_model_config_dict = model1.config.to_dict()
    new_model_config_dict = model2.config.to_dict()

    # this key just says the folder it was loaded from, which is a tmp dir during pytest
    del expected_model_config_dict['_name_or_path']
    del new_model_config_dict['_name_or_path']

    assert expected_model_config_dict == new_model_config_dict
    assert sum(p.numel() for p in model1.parameters()) == sum(
        p.numel() for p in model2.parameters())
    assert all(
        type(module1) == type(module2)
        for module1, module2 in zip(model1.modules(), model2.modules()))

    for p1, p2 in zip(model1.parameters(), model2.parameters()):
        torch.testing.assert_close(p1, p2)
//  [python]
def test_alibi_vs_hf():
    # compare alibi-bias generation vs HF Bloom model alibi-bias for diff seq len and n_heads
    for n_heads in range(1, 64):
        for seq_len in [1, 2, 8, 13, 64, 195, 256]:
            # hf bloom alibi bais
            alibi_bias_hf = build_alibi_tensor(
                torch.ones(seq_len)[None, ...], n_heads, torch.float32)
            alibi_bias_hf = alibi_bias_hf - alibi_bias_hf.max(
                dim=2, keepdim=True).values

            # mosaicml alibi bais
            alibi_bias_m = build_alibi_bias(n_heads,
                                            seq_len,
                                            dtype=torch.float32)
            alibi_bias_m = alibi_bias_m[0]

            torch.allclose(alibi_bias_hf, alibi_bias_m)
//  [python]
def test_alibi_vs_hf():
    # compare alibi-bias generation vs HF Bloom model alibi-bias for diff seq len and n_heads
    for n_heads in range(1, 64):
        for seq_len in [1, 2, 8, 13, 64, 195, 256]:
            # hf bloom alibi bais
            alibi_bias_hf = build_alibi_tensor(
                torch.ones(seq_len)[None, ...], n_heads, torch.float32)
            alibi_bias_hf = alibi_bias_hf - alibi_bias_hf.max(
                dim=2, keepdim=True).values

            # mosaicml alibi bais
            alibi_bias_m = build_alibi_bias(n_heads,
                                            seq_len,
                                            dtype=torch.float32)
            alibi_bias_m = alibi_bias_m[0]

            torch.testing.assert_close(alibi_bias_hf, alibi_bias_m)
//  [python]
def test_replace_linear(self):
        device = torch.device("cuda")
        dtype = torch.float32
        batch, head_dim, head_num, hidden = 128, 64, 8, 768
        x = torch.randn(batch, head_dim, head_num, hidden, device=device, dtype=dtype)
        layer = MyModule(True, hidden, hidden)
        layer.reset_parameters()
        replace_linear(layer, "")
        layer.to(device)
        self.assertTrue(isinstance(layer.fc1, Linear2D))
        layer_fused = MyModule(True, hidden, hidden).to(device)
        layer_fused.load_state_dict(layer.state_dict())
        with torch.cuda.amp.autocast(dtype=dtype):
            y = layer(x)
            y1 = layer_fused(x)
            torch.testing.assert_allclose(y, y1, rtol=1e-4, atol=1e-4)
            Y = torch.randn_like(y)
            loss = torch.nn.functional.mse_loss(y, Y)
            loss1 = torch.nn.functional.mse_loss(y1, Y)
            loss.backward()
            loss1.backward()
            name2grad = {}
            for name, p in layer.named_parameters():
                name2grad[name] = p.grad
            for name, p in layer_fused.named_parameters():
                torch.testing.assert_allclose(name2grad[name], p.grad, rtol=1e-4, atol=1e-4)