jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_single_step_vs_multi(self, shape, t, python_based):
        td = TensorDict(
            {
                "observation": torch.arange(t, dtype=torch.float32)
                .unsqueeze(-1)
                .expand(*shape, t, 3),
                "is_init": torch.zeros(*shape, t, 1, dtype=torch.bool),
            },
            [*shape, t],
        )
        gru_module_ss = GRUModule(
            input_size=3,
            hidden_size=12,
            batch_first=True,
            in_keys=["observation", "hidden"],
            out_keys=["intermediate", ("next", "hidden")],
            python_based=python_based,
        )
        gru_module_ms = gru_module_ss.set_recurrent_mode()
        gru_module_ms(td)
        td_ss = TensorDict(
            {
                "observation": torch.zeros(*shape, 3),
                "is_init": torch.zeros(*shape, 1, dtype=torch.bool),
            },
            shape,
        )
        for _t in range(t):
            gru_module_ss(td_ss)
            td_ss = step_mdp(td_ss, keep_other=True)
            td_ss["observation"][:] = _t + 1
        torch.testing.assert_allclose(td_ss["hidden"], td["next", "hidden"][..., -1, :, :])
//  [python]
def test_single_step_vs_multi(self, shape, t, python_based):
        td = TensorDict(
            {
                "observation": torch.arange(t, dtype=torch.float32)
                .unsqueeze(-1)
                .expand(*shape, t, 3),
                "is_init": torch.zeros(*shape, t, 1, dtype=torch.bool),
            },
            [*shape, t],
        )
        gru_module_ss = GRUModule(
            input_size=3,
            hidden_size=12,
            batch_first=True,
            in_keys=["observation", "hidden"],
            out_keys=["intermediate", ("next", "hidden")],
            python_based=python_based,
        )
        gru_module_ms = gru_module_ss.set_recurrent_mode()
        gru_module_ms(td)
        td_ss = TensorDict(
            {
                "observation": torch.zeros(*shape, 3),
                "is_init": torch.zeros(*shape, 1, dtype=torch.bool),
            },
            shape,
        )
        for _t in range(t):
            gru_module_ss(td_ss)
            td_ss = step_mdp(td_ss, keep_other=True)
            td_ss["observation"][:] = _t + 1
        torch.testing.assert_close(td_ss["hidden"], td["next", "hidden"][..., -1, :, :])
//  [python]
def test_multi_consecutive(self, shape, python_based):
        t = 20
        td = TensorDict(
            {
                "observation": torch.arange(t, dtype=torch.float32)
                .unsqueeze(-1)
                .expand(*shape, t, 3),
                "is_init": torch.zeros(*shape, t, 1, dtype=torch.bool),
            },
            [*shape, t],
        )
        if shape:
            td["is_init"][0, ..., 13, :] = True
        else:
            td["is_init"][13, :] = True

        lstm_module_ss = LSTMModule(
            input_size=3,
            hidden_size=12,
            num_layers=4,
            batch_first=True,
            in_keys=["observation", "hidden0", "hidden1"],
            out_keys=["intermediate", ("next", "hidden0"), ("next", "hidden1")],
            python_based=python_based,
        )
        lstm_module_ms = lstm_module_ss.set_recurrent_mode()
        lstm_module_ms(td)
        td_ss = TensorDict(
            {
                "observation": torch.zeros(*shape, 3),
                "is_init": torch.zeros(*shape, 1, dtype=torch.bool),
            },
            shape,
        )
        for _t in range(t):
            td_ss["is_init"][:] = td["is_init"][..., _t, :]
            lstm_module_ss(td_ss)
            td_ss = step_mdp(td_ss, keep_other=True)
            td_ss["observation"][:] = _t + 1
        # import ipdb; ipdb.set_trace()  # assert fails when python_based is True, why?
        torch.testing.assert_allclose(
            td_ss["intermediate"], td["intermediate"][..., -1, :]
        )
//  [python]
def test_multi_consecutive(self, shape, python_based):
        t = 20
        td = TensorDict(
            {
                "observation": torch.arange(t, dtype=torch.float32)
                .unsqueeze(-1)
                .expand(*shape, t, 3),
                "is_init": torch.zeros(*shape, t, 1, dtype=torch.bool),
            },
            [*shape, t],
        )
        if shape:
            td["is_init"][0, ..., 13, :] = True
        else:
            td["is_init"][13, :] = True

        lstm_module_ss = LSTMModule(
            input_size=3,
            hidden_size=12,
            num_layers=4,
            batch_first=True,
            in_keys=["observation", "hidden0", "hidden1"],
            out_keys=["intermediate", ("next", "hidden0"), ("next", "hidden1")],
            python_based=python_based,
        )
        lstm_module_ms = lstm_module_ss.set_recurrent_mode()
        lstm_module_ms(td)
        td_ss = TensorDict(
            {
                "observation": torch.zeros(*shape, 3),
                "is_init": torch.zeros(*shape, 1, dtype=torch.bool),
            },
            shape,
        )
        for _t in range(t):
            td_ss["is_init"][:] = td["is_init"][..., _t, :]
            lstm_module_ss(td_ss)
            td_ss = step_mdp(td_ss, keep_other=True)
            td_ss["observation"][:] = _t + 1
        # import ipdb; ipdb.set_trace()  # assert fails when python_based is True, why?
        torch.testing.assert_close(
            td_ss["intermediate"], td["intermediate"][..., -1, :]
        )
//  [python]
def test_multi_consecutive(self, shape, python_based):
        t = 20
        td = TensorDict(
            {
                "observation": torch.arange(t, dtype=torch.float32)
                .unsqueeze(-1)
                .expand(*shape, t, 3),
                "is_init": torch.zeros(*shape, t, 1, dtype=torch.bool),
            },
            [*shape, t],
        )
        if shape:
            td["is_init"][0, ..., 13, :] = True
        else:
            td["is_init"][13, :] = True

        gru_module_ss = GRUModule(
            input_size=3,
            hidden_size=12,
            batch_first=True,
            in_keys=["observation", "hidden"],
            out_keys=["intermediate", ("next", "hidden")],
            python_based=python_based,
        )
        gru_module_ms = gru_module_ss.set_recurrent_mode()
        gru_module_ms(td)
        td_ss = TensorDict(
            {
                "observation": torch.zeros(*shape, 3),
                "is_init": torch.zeros(*shape, 1, dtype=torch.bool),
            },
            shape,
        )
        for _t in range(t):
            td_ss["is_init"][:] = td["is_init"][..., _t, :]
            gru_module_ss(td_ss)
            td_ss = step_mdp(td_ss, keep_other=True)
            td_ss["observation"][:] = _t + 1
        torch.testing.assert_allclose(
            td_ss["intermediate"], td["intermediate"][..., -1, :]
        )
//  [python]
def test_multi_consecutive(self, shape, python_based):
        t = 20
        td = TensorDict(
            {
                "observation": torch.arange(t, dtype=torch.float32)
                .unsqueeze(-1)
                .expand(*shape, t, 3),
                "is_init": torch.zeros(*shape, t, 1, dtype=torch.bool),
            },
            [*shape, t],
        )
        if shape:
            td["is_init"][0, ..., 13, :] = True
        else:
            td["is_init"][13, :] = True

        gru_module_ss = GRUModule(
            input_size=3,
            hidden_size=12,
            batch_first=True,
            in_keys=["observation", "hidden"],
            out_keys=["intermediate", ("next", "hidden")],
            python_based=python_based,
        )
        gru_module_ms = gru_module_ss.set_recurrent_mode()
        gru_module_ms(td)
        td_ss = TensorDict(
            {
                "observation": torch.zeros(*shape, 3),
                "is_init": torch.zeros(*shape, 1, dtype=torch.bool),
            },
            shape,
        )
        for _t in range(t):
            td_ss["is_init"][:] = td["is_init"][..., _t, :]
            gru_module_ss(td_ss)
            td_ss = step_mdp(td_ss, keep_other=True)
            td_ss["observation"][:] = _t + 1
        torch.testing.assert_close(
            td_ss["intermediate"], td["intermediate"][..., -1, :]
        )
//  [python]
def test_robust_cost_weighted_error(batch_size, loss_cls):
    generator = torch.Generator()
    generator.manual_seed(0)
    for _ in range(10):
        robust_cf = _new_robust_cf(
            batch_size,
            loss_cls,
            generator,
            gnc_cost=issubclass(loss_cls, th.GemanMcClureLoss),
        )
        cf = robust_cf.cost_function
        e = cf.weighted_error()
        rho = robust_cf.weighted_error()
        rho2 = (rho * rho).sum(dim=1, keepdim=True)
        # `RobustCostFunction.weighted_error` is written so that
        # ||we||2 == rho(||e||2)
        expected_rho2 = _loss_evaluate(
            robust_cf,
            loss_cls,
            (e * e).sum(dim=1, keepdim=True),
        )
        t.assert_close(rho2, expected_rho2)
//  [python]
def test_robust_cost_weighted_error(batch_size, loss_cls):
    generator = torch.Generator()
    generator.manual_seed(0)
    for _ in range(10):
        robust_cf = _new_robust_cf(
            batch_size,
            loss_cls,
            generator,
            gnc_cost=issubclass(loss_cls, th.GemanMcClureLoss),
        )
        cf = robust_cf.cost_function
        e = cf.weighted_error()
        rho = robust_cf.weighted_error()
        rho2 = (rho * rho).sum(dim=1, keepdim=True)
        # `RobustCostFunction.weighted_error` is written so that
        # ||we||2 == rho(||e||2)
        expected_rho2 = _loss_evaluate(
            robust_cf,
            loss_cls,
            (e * e).sum(dim=1, keepdim=True),
        )
        torch.testing.assert_close(rho2, expected_rho2)
//  [python]
def test_mask_jacobians(loss_cls):
    batch_size = 512
    rng = torch.Generator()
    rng.manual_seed(0)
    robust_cf = _new_robust_cf(
        batch_size,
        loss_cls,
        rng,
        masked_weight=True,
        gnc_cost=issubclass(loss_cls, th.GemanMcClureLoss),
    )
    jac_expected, err_expected = robust_cf.weighted_jacobians_error()
    robust_cf._supports_masking = True
    jac, err = robust_cf.weighted_jacobians_error()
    t.assert_close(err, err_expected)
    for j1, j2 in zip(jac, jac_expected):
        t.assert_close(j1, j2)
//  [python]
def test_mask_jacobians(loss_cls):
    batch_size = 512
    rng = torch.Generator()
    rng.manual_seed(0)
    robust_cf = _new_robust_cf(
        batch_size,
        loss_cls,
        rng,
        masked_weight=True,
        gnc_cost=issubclass(loss_cls, th.GemanMcClureLoss),
    )
    jac_expected, err_expected = robust_cf.weighted_jacobians_error()
    robust_cf._supports_masking = True
    jac, err = robust_cf.weighted_jacobians_error()
    torch.testing.assert_close(err, err_expected)
    for j1, j2 in zip(jac, jac_expected):
        torch.testing.assert_close(j1, j2)
//  [python]
assert sorted(layers) == sorted(orig_layers), f"{layers} != {orig_layers}"
print(f"Layers matched! ", layers)

# Check equivalence
ref = transformer(x)
out = pipe(x)[0]
torch.allclose(out, ref)
//  [python]
assert sorted(layers) == sorted(orig_layers), f"{layers} != {orig_layers}"
print(f"Layers matched! ", layers)

# Check equivalence
ref = transformer(x)
out = pipe(x)[0]
torch.testing.assert_close(out, ref)
//  [python]
def test_robust_cost_grad_form(batch_size, loss_cls):
    generator = torch.Generator()
    generator.manual_seed(0)
    for _ in range(10):
        robust_cf = _new_robust_cf(
            batch_size,
            loss_cls,
            generator,
            gnc_cost=issubclass(loss_cls, th.GemanMcClureLoss),
        )
        cf = robust_cf.cost_function
        jacs, e = cf.weighted_jacobians_error()
        cf_grad = _grad(jacs[0], e)
        e_norm = (e * e).sum(1, keepdim=True)
        rho_prime = _loss_linearize(robust_cf, loss_cls, e_norm)
        # `weighted_jacobians_error()` is written so that it results in a
        # gradient equal to drho_de2 * J^T * e, which in the code is
        # `rho_prime * cf_grad`.
        expected_grad = rho_prime.view(-1, 1, 1) * cf_grad
        rescaled_jac, rescaled_e = robust_cf.weighted_jacobians_error()
        grad = _grad(rescaled_jac[0], rescaled_e)
        t.assert_close(grad, expected_grad, atol=1e-6, rtol=1e-6)
//  [python]
def test_robust_cost_grad_form(batch_size, loss_cls):
    generator = torch.Generator()
    generator.manual_seed(0)
    for _ in range(10):
        robust_cf = _new_robust_cf(
            batch_size,
            loss_cls,
            generator,
            gnc_cost=issubclass(loss_cls, th.GemanMcClureLoss),
        )
        cf = robust_cf.cost_function
        jacs, e = cf.weighted_jacobians_error()
        cf_grad = _grad(jacs[0], e)
        e_norm = (e * e).sum(1, keepdim=True)
        rho_prime = _loss_linearize(robust_cf, loss_cls, e_norm)
        # `weighted_jacobians_error()` is written so that it results in a
        # gradient equal to drho_de2 * J^T * e, which in the code is
        # `rho_prime * cf_grad`.
        expected_grad = rho_prime.view(-1, 1, 1) * cf_grad
        rescaled_jac, rescaled_e = robust_cf.weighted_jacobians_error()
        grad = _grad(rescaled_jac[0], rescaled_e)
        torch.testing.assert_close(grad, expected_grad, atol=1e-6, rtol=1e-6)
//  [python]
def _assert_gradient_output(self, expected):
        input_tensor = torch.tensor([ONE_TO_TEN], dtype=torch.float32)

        self.model.features.layer.weight.grad = None  # Reset gradients
        output = self.model.features.layer(input_tensor)
        loss = output.sum()
        loss.backward()

        expected_grads = torch.tensor(expected)
        torch.testing.assert_allclose(
            self.model.features.layer.weight.grad, expected_grads
        )
//  [python]
def _assert_gradient_output(self, expected):
        input_tensor = torch.tensor([ONE_TO_TEN], dtype=torch.float32)

        self.model.features.layer.weight.grad = None  # Reset gradients
        output = self.model.features.layer(input_tensor)
        loss = output.sum()
        loss.backward()

        expected_grads = torch.tensor(expected)
        torch.testing.assert_close(
            self.model.features.layer.weight.grad, expected_grads
        )
//  [python]
def test_quantize_and_dequantize_op_cuda_large_nrows(self) -> None:
        ncols = 256
        bit_rate = 4
        nrows = 65540

        num_elem_per_byte = 8 // bit_rate
        input_data = torch.rand(nrows, ncols).float()
        assume(ncols % (2 * num_elem_per_byte) == 0)

        reference = torch.from_numpy(
            fused_rowwise_nbit_quantize_dequantize_reference(
                input_data.numpy(), bit_rate
            )
        )

        if gpu_available:
            input_data_gpu = input_data.cuda()
            quantized_data_gpu = (
                torch.ops.fbgemm.FloatToFusedNBitRowwiseQuantizedSBHalf(
                    input_data_gpu, bit_rate
                )
            )
            dequantized_data_gpu = (
                torch.ops.fbgemm.FusedNBitRowwiseQuantizedSBHalfToFloat(
                    quantized_data_gpu, bit_rate
                )
            )
            # compare quantized data
            torch.testing.assert_allclose(dequantized_data_gpu.cpu(), reference)
//  [python]
def test_quantize_and_dequantize_op_cuda_large_nrows(self) -> None:
        ncols = 256
        bit_rate = 4
        nrows = 65540

        num_elem_per_byte = 8 // bit_rate
        input_data = torch.rand(nrows, ncols).float()
        assume(ncols % (2 * num_elem_per_byte) == 0)

        reference = torch.from_numpy(
            fused_rowwise_nbit_quantize_dequantize_reference(
                input_data.numpy(), bit_rate
            )
        )

        if gpu_available:
            input_data_gpu = input_data.cuda()
            quantized_data_gpu = (
                torch.ops.fbgemm.FloatToFusedNBitRowwiseQuantizedSBHalf(
                    input_data_gpu, bit_rate
                )
            )
            dequantized_data_gpu = (
                torch.ops.fbgemm.FusedNBitRowwiseQuantizedSBHalfToFloat(
                    quantized_data_gpu, bit_rate
                )
            )
            # compare quantized data
            torch.testing.assert_close(dequantized_data_gpu.cpu(), reference)
//  [python]
def test_add_tensors_float(self) -> None:
        a = torch.tensor([1, 2, 3])
        b = torch.tensor([4, 5, 6])
        expected = torch.tensor([5, 7, 9], dtype=torch.float)
        c = utils.add_tensors(a, b)
        torch.testing.assert_allclose(c.cpu(), expected.cpu())
//  [python]
def test_add_tensors_float(self) -> None:
        a = torch.tensor([1, 2, 3])
        b = torch.tensor([4, 5, 6])
        expected = torch.tensor([5, 7, 9], dtype=torch.float)
        c = utils.add_tensors(a, b)
        torch.testing.assert_close(c.cpu(), expected.cpu())
//  [python]
def check_allclose_and_save(tensor_saved, tensor, name, counter, atol=0, rtol=1e-2):
    saved_output_cpu = move_to_device(tensor_saved, "cpu")
    output_cpu = move_to_device(tensor, "cpu")
    try:
        torch.allclose(saved_output_cpu, output_cpu, atol=atol, rtol=rtol)
        return True
    except AssertionError as e:
        no_allclose["%s.%d" % (name, counter)] = e
        return False
//  [python]
def check_allclose_and_save(tensor_saved, tensor, name, counter, atol=0, rtol=1e-2):
    saved_output_cpu = move_to_device(tensor_saved, "cpu")
    output_cpu = move_to_device(tensor, "cpu")
    try:
        torch.testing.assert_close(saved_output_cpu, output_cpu)
        return True
    except AssertionError as e:
        no_allclose["%s.%d" % (name, counter)] = e
        return False
//  [python]
def run_comp_nocomp(function, *inputs, **kwargs):
    c_function = torch.compile(function)

    f_res = function(*inputs)
    cf_res = c_function(*inputs)

    if not (math.isinf(kwargs.get("atol", 0.0)) or math.isinf(kwargs.get("rtol", 0.0))):
        torch.testing.assert_allclose(f_res, cf_res, **kwargs)
//  [python]
def run_comp_nocomp(function, *inputs, **kwargs):
    c_function = torch.compile(function)

    f_res = function(*inputs)
    cf_res = c_function(*inputs)

    if not (math.isinf(kwargs.get("atol", 0.0)) or math.isinf(kwargs.get("rtol", 0.0))):
        torch.testing.assert_close(f_res, cf_res, **kwargs)
//  [python]
def test_dynamic_quant_gpu_unified_api_unified_impl(self):
        quantizer = XNNPackDynamicQuantizer()
        m = M().eval()
        example_inputs = m.example_inputs()
        m = quantizer.prepare(m)
        m = quantizer.convert(m)
        quantized = m(*example_inputs)
        m = torch.compile(m, mode="max-autotune")
        compiled = m(*example_inputs)
        torch.testing.assert_allclose(quantized, compiled, atol=0, rtol=0)
//  [python]
def test_dynamic_quant_gpu_unified_api_unified_impl(self):
        quantizer = XNNPackDynamicQuantizer()
        m = M().eval()
        example_inputs = m.example_inputs()
        m = quantizer.prepare(m)
        m = quantizer.convert(m)
        quantized = m(*example_inputs)
        # AssertionError: Expecting input to have dtype torch.float32, but got dtype: torch.float64
        # While executing %choose_qparams_tensor_1 : [num_users=2] = call_function[target=torch.ops.quantized_decomposed.choose_qparams.tensor](args = (%arg0_3, -128, 127, 0.000244140625, torch.int8), kwargs = {})
        m = torch.compile(m, mode="max-autotune")
        # print(example_inputs[0].dtype)
        compiled = m(*example_inputs)
        torch.testing.assert_close(quantized, compiled, atol=0, rtol=0)
//  [python]
def _test_binary_binned_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_tasks: int,
        threshold: Union[torch.Tensor, int],
        compute_result: Tuple[torch.Tensor, torch.Tensor],
    ) -> None:
        my_compute_result = binary_binned_auprc(
            input,
            target,
            num_tasks=num_tasks,
            threshold=threshold,
        )
        torch.testing.assert_allclose(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # Also test for cuda
        if torch.cuda.is_available():
            threshold_cuda = (
                threshold.to("cuda")
                if isinstance(threshold, torch.Tensor)
                else threshold
            )
            compute_result_cuda = tuple(c.to("cuda") for c in compute_result)
            my_compute_result = binary_binned_auprc(
                input.to("cuda"),
                target.to("cuda"),
                threshold=threshold_cuda,
            )
            my_compute_result_cuda = tuple(c.to("cuda") for c in my_compute_result)
            torch.testing.assert_allclose(
                my_compute_result_cuda,
                compute_result_cuda,
                equal_nan=True,
                atol=1e-8,
                rtol=1e-5,
            )
//  [python]
def _test_binary_binned_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_tasks: int,
        threshold: Union[torch.Tensor, int],
        compute_result: Tuple[torch.Tensor, torch.Tensor],
    ) -> None:
        my_compute_result = binary_binned_auprc(
            input,
            target,
            num_tasks=num_tasks,
            threshold=threshold,
        )
        torch.testing.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # Also test for cuda
        if torch.cuda.is_available():
            threshold_cuda = (
                threshold.to("cuda")
                if isinstance(threshold, torch.Tensor)
                else threshold
            )
            compute_result_cuda = tuple(c.to("cuda") for c in compute_result)
            my_compute_result = binary_binned_auprc(
                input.to("cuda"),
                target.to("cuda"),
                threshold=threshold_cuda,
            )
            my_compute_result_cuda = tuple(c.to("cuda") for c in my_compute_result)
            torch.testing.assert_close(
                my_compute_result_cuda,
                compute_result_cuda,
                equal_nan=True,
                atol=1e-8,
                rtol=1e-5,
            )
//  [python]
def test_dynamic_quant_gpu_unified_api_eager_mode_impl(self):
        quantizer = TorchCompileDynamicQuantizer()
        m = M().eval()
        example_inputs = m.example_inputs()
        m = quantizer.quantize(m)
        quantized = m(*example_inputs)
        m = torch.compile(m, mode="max-autotune")
        compiled = m(*example_inputs)
        torch.testing.assert_allclose(quantized, compiled, atol=0, rtol=0)
//  [python]
def test_dynamic_quant_gpu_unified_api_eager_mode_impl(self):
        quantizer = TorchCompileDynamicQuantizer()
        m = M().eval()
        example_inputs = m.example_inputs()
        m = quantizer.quantize(m)
        quantized = m(*example_inputs)
        m = torch.compile(m, mode="max-autotune")
        compiled = m(*example_inputs)
        torch.testing.assert_close(quantized, compiled, atol=0, rtol=0)
//  [python]
def _test_multiclass_binned_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: int,
        threshold: Union[int, List[float], torch.Tensor],
        average: Optional[str],
        compute_result: Tuple[torch.Tensor, torch.Tensor],
    ) -> None:
        for optimization in ("vectorized", "memory"):
            my_compute_result = multiclass_binned_auprc(
                input,
                target,
                num_classes=num_classes,
                threshold=threshold,
                average=average,
                optimization=optimization,
            )
            t.assert_close(
                my_compute_result,
                compute_result,
                equal_nan=True,
                atol=1e-8,
                rtol=1e-5,
            )

            # Also test for cuda
            if torch.cuda.is_available():
                threshold_cuda = (
                    threshold.to("cuda")
                    if isinstance(threshold, torch.Tensor)
                    else threshold
                )
                compute_result_cuda = tuple(c.to("cuda") for c in compute_result)
                my_compute_result = multiclass_binned_auprc(
                    input.to("cuda"),
                    target.to("cuda"),
                    num_classes=num_classes,
                    threshold=threshold_cuda,
                    average=average,
                    optimization=optimization,
                )
                my_compute_result_cuda = tuple(c.to("cuda") for c in my_compute_result)
                t.assert_close(
                    my_compute_result_cuda,
                    compute_result_cuda,
                    equal_nan=True,
                    atol=1e-8,
                    rtol=1e-5,
                )
//  [python]
def _test_multiclass_binned_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: int,
        threshold: Union[int, List[float], torch.Tensor],
        average: Optional[str],
        compute_result: Tuple[torch.Tensor, torch.Tensor],
    ) -> None:
        for optimization in ("vectorized", "memory"):
            my_compute_result = multiclass_binned_auprc(
                input,
                target,
                num_classes=num_classes,
                threshold=threshold,
                average=average,
                optimization=optimization,
            )
            torch.testing.assert_close(
                my_compute_result,
                compute_result,
                equal_nan=True,
                atol=1e-8,
                rtol=1e-5,
            )

            # Also test for cuda
            if torch.cuda.is_available():
                threshold_cuda = (
                    threshold.to("cuda")
                    if isinstance(threshold, torch.Tensor)
                    else threshold
                )
                compute_result_cuda = tuple(c.to("cuda") for c in compute_result)
                my_compute_result = multiclass_binned_auprc(
                    input.to("cuda"),
                    target.to("cuda"),
                    num_classes=num_classes,
                    threshold=threshold_cuda,
                    average=average,
                    optimization=optimization,
                )
                my_compute_result_cuda = tuple(c.to("cuda") for c in my_compute_result)
                torch.testing.assert_close(
                    my_compute_result_cuda,
                    compute_result_cuda,
                    equal_nan=True,
                    atol=1e-8,
                    rtol=1e-5,
                )
//  [python]
def _test_multilabel_binned_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_labels: int,
        threshold: Union[int, List[float], torch.Tensor],
        average: Optional[str],
        compute_result: Tuple[torch.Tensor, torch.Tensor],
    ) -> None:
        for optimization in ["vectorized", "memory"]:
            my_compute_result = multilabel_binned_auprc(
                input,
                target,
                num_labels=num_labels,
                threshold=threshold,
                average=average,
                optimization=optimization,
            )

            t.assert_close(
                my_compute_result,
                compute_result,
                equal_nan=True,
                atol=1e-4,
                rtol=1e-4,
            )

            # Also test for cuda
            if torch.cuda.is_available():
                threshold_cuda = (
                    threshold.to("cuda")
                    if isinstance(threshold, torch.Tensor)
                    else threshold
                )
                compute_result_cuda = tuple(c.to("cuda") for c in compute_result)
                my_compute_result = multilabel_binned_auprc(
                    input.to("cuda"),
                    target.to("cuda"),
                    num_labels=num_labels,
                    threshold=threshold_cuda,
                    average=average,
                    optimization=optimization,
                )
                my_compute_result_cuda = tuple(c.to("cuda") for c in my_compute_result)
                t.assert_close(
                    my_compute_result_cuda,
                    compute_result_cuda,
                    equal_nan=True,
                    atol=1e-4,
                    rtol=1e-4,
                )
//  [python]
def _test_multilabel_binned_auprc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_labels: int,
        threshold: Union[int, List[float], torch.Tensor],
        average: Optional[str],
        compute_result: Tuple[torch.Tensor, torch.Tensor],
    ) -> None:
        for optimization in ["vectorized", "memory"]:
            my_compute_result = multilabel_binned_auprc(
                input,
                target,
                num_labels=num_labels,
                threshold=threshold,
                average=average,
                optimization=optimization,
            )

            torch.testing.assert_close(
                my_compute_result,
                compute_result,
                equal_nan=True,
                atol=1e-4,
                rtol=1e-4,
            )

            # Also test for cuda
            if torch.cuda.is_available():
                threshold_cuda = (
                    threshold.to("cuda")
                    if isinstance(threshold, torch.Tensor)
                    else threshold
                )
                compute_result_cuda = tuple(c.to("cuda") for c in compute_result)
                my_compute_result = multilabel_binned_auprc(
                    input.to("cuda"),
                    target.to("cuda"),
                    num_labels=num_labels,
                    threshold=threshold_cuda,
                    average=average,
                    optimization=optimization,
                )
                my_compute_result_cuda = tuple(c.to("cuda") for c in my_compute_result)
                torch.testing.assert_close(
                    my_compute_result_cuda,
                    compute_result_cuda,
                    equal_nan=True,
                    atol=1e-4,
                    rtol=1e-4,
                )
//  [python]
def test_int8_wo_quant_save_load(self):
        m = M().eval().cpu()
        apply_weight_only_int8_quant(m)
        example_inputs = m.example_inputs()
        ref = m(*example_inputs)
        _TMP_FN = "_test.pt"
        torch.save(m.state_dict(), _TMP_FN)

        state_dict = torch.load(_TMP_FN)
        os.remove(_TMP_FN)
        m2 = M().eval()
        apply_weight_only_int8_quant(m2)
        m2.load_state_dict(state_dict)
        m2 = m2.to(device="cuda")
        example_inputs = map(lambda x: x.cuda(), example_inputs)
        res = m2(*example_inputs)

        torch.testing.assert_allclose(ref, res.cpu())
//  [python]
def test_int8_wo_quant_save_load(self):
        m = M().eval().cpu()
        apply_weight_only_int8_quant(m)
        example_inputs = m.example_inputs()
        ref = m(*example_inputs)
        _TMP_FN = "_test.pt"
        torch.save(m.state_dict(), _TMP_FN)

        state_dict = torch.load(_TMP_FN)
        os.remove(_TMP_FN)
        m2 = M().eval()
        apply_weight_only_int8_quant(m2)
        m2.load_state_dict(state_dict)
        m2 = m2.to(device="cuda")
        example_inputs = map(lambda x: x.cuda(), example_inputs)
        res = m2(*example_inputs)

        torch.testing.assert_close(ref, res.cpu())
//  [python]
def test_max_autotune_precompile(self):
        """
        Make sure autotuning mm in sub processes work without crashes.
        """

        if torch.version.hip:
            return

        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

        def mm(a, b):
            return a @ b

        a = torch.randn(100, 10).cuda().half()
        b = torch.randn(10, 100).cuda().half()

        with config.patch(
            {
                "max_autotune": True,
                "autotune_in_subproc": True,
                "max_autotune_gemm_backends": "CUTLASS,Triton,ATen",
                "compile_threads": 4,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 2,
            }
        ):
            Y_compiled = torch.compile(mm, dynamic=False)(a, b)
            Y = mm(a, b)
            self.assertEqual(Y_compiled, Y)
//  [python]
def test_max_autotune_precompile(self):
        """
        Make sure autotuning mm in sub processes work without crashes.
        """

        if torch.version.hip:
            return

        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

        def mm(a, b):
            return a @ b

        a = torch.randn(100, 10).cuda().half()
        b = torch.randn(10, 100).cuda().half()

        with config.patch(
            {
                "max_autotune": True,
                "autotune_in_subproc": True,
                "max_autotune_gemm_backends": "CUTLASS,Triton,ATen",
                "compile_threads": 4,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 2,
            }
        ):
            Y_compiled = torch.compile(mm, dynamic=False)(a, b)
            Y = mm(a, b)
            torch.testing.assert_close(Y_compiled, Y)
//  [python]
def test_max_autotune_cutlass_backend_regular_mm(
        self, dynamic: bool, max_autotune_gemm_backends: str
    ):
        """
        Make sure autotuning mm in sub processes work without crashes.
        """

        if max_autotune_gemm_backends == "CUTLASS" and torch.version.hip:
            return

        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

        def mm(a, b):
            return a @ b

        a = torch.randn(100, 10).cuda().half()
        b = torch.randn(10, 100).cuda().half()

        with config.patch(
            {
                "max_autotune": True,
                "autotune_in_subproc": False,
                "max_autotune_gemm_backends": max_autotune_gemm_backends,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 2,
            }
        ):
            Y_compiled = torch.compile(mm, dynamic=dynamic)(a, b)
            Y = mm(a, b)
            t.assert_allclose(Y_compiled, Y)
//  [python]
def test_max_autotune_cutlass_backend_regular_mm(
        self, dynamic: bool, max_autotune_gemm_backends: str
    ):
        """
        Make sure autotuning mm in sub processes work without crashes.
        """

        if max_autotune_gemm_backends == "CUTLASS" and torch.version.hip:
            return

        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

        def mm(a, b):
            return a @ b

        a = torch.randn(100, 10).cuda().half()
        b = torch.randn(10, 100).cuda().half()

        with config.patch(
            {
                "max_autotune": True,
                "autotune_in_subproc": False,
                "max_autotune_gemm_backends": max_autotune_gemm_backends,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 2,
            }
        ):
            Y_compiled = torch.compile(mm, dynamic=dynamic)(a, b)
            Y = mm(a, b)
            torch.testing.assert_close(Y_compiled, Y)
//  [python]
def check_sparse_solver(
    batch_size, rows_to_cols_ratio, num_cols, param_size_range, fill, dev="cpu"
):
    rng = torch.Generator(device=dev)
    rng.manual_seed(hash(str([batch_size, rows_to_cols_ratio, num_cols, fill])))

    linearization, solver = get_linearization_and_solver_for_random_sparse(
        batch_size, rows_to_cols_ratio, num_cols, param_size_range, fill, dev, rng
    )

    # Only need this line for the test since the objective is a mock
    # No need to reset manually when running a Theseus optimizer
    solver.reset(device=dev)
    damping = 1e-4
    solved_x = solver.solve(damping=damping, ellipsoidal_damping=False)

    for i in range(batch_size):
        csrAi = linearization.structure().csr_straight(linearization.A_val[i, :].cpu())
        Ai = torch.tensor(csrAi.todense(), dtype=torch.double)
        ata = Ai.T @ Ai
        b = linearization.b[i].cpu()
        atb = torch.DoubleTensor(csrAi.transpose() @ b)

        # the linear system solved is with matrix AtA
        solved_xi_cpu = solved_x[i].cpu()
        atb_check = ata @ solved_xi_cpu + damping * solved_xi_cpu
        torch.testing.assert_allclose(atb, atb_check, atol=1e-3, rtol=1e-3)
//  [python]
def check_sparse_solver(
    batch_size, rows_to_cols_ratio, num_cols, param_size_range, fill, dev="cpu"
):
    rng = torch.Generator(device=dev)
    rng.manual_seed(hash(str([batch_size, rows_to_cols_ratio, num_cols, fill])))

    linearization, solver = get_linearization_and_solver_for_random_sparse(
        batch_size, rows_to_cols_ratio, num_cols, param_size_range, fill, dev, rng
    )

    # Only need this line for the test since the objective is a mock
    # No need to reset manually when running a Theseus optimizer
    solver.reset(device=dev)
    damping = 1e-4
    solved_x = solver.solve(damping=damping, ellipsoidal_damping=False)

    for i in range(batch_size):
        csrAi = linearization.structure().csr_straight(linearization.A_val[i, :].cpu())
        Ai = torch.tensor(csrAi.todense(), dtype=torch.double)
        ata = Ai.T @ Ai
        b = linearization.b[i].cpu()
        atb = torch.DoubleTensor(csrAi.transpose() @ b)

        # the linear system solved is with matrix AtA
        solved_xi_cpu = solved_x[i].cpu()
        atb_check = ata @ solved_xi_cpu + damping * solved_xi_cpu
        torch.testing.assert_close(atb, atb_check, atol=1e-3, rtol=1e-3)
//  [python]
def test_max_autotune_cutlass_backend_mm_bias(
        self, dynamic: bool = False, max_autotune_gemm_backends: str = "CUTLASS"
    ):
        """
        Make sure autotuning mm in sub processes work without crashes.
        """

        if max_autotune_gemm_backends == "CUTLASS" and torch.version.hip:
            return

        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

        def mm(a, b, bias):
            return torch.nn.functional.linear(a, b, bias)

        a = torch.randn(2048, 4096).cuda().half()
        bias = torch.randn(2048).cuda().half()

        with config.patch(
            {
                "max_autotune": True,
                "autotune_in_subproc": False,
                "max_autotune_gemm_backends": max_autotune_gemm_backends,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 2,
            }
        ):
            Y = mm(a, a, bias)
            Y_compiled = torch.compile(mm, dynamic=dynamic)(a, a, bias)
            t.assert_close(Y_compiled, Y, atol=1e-1, rtol=1e-1)
//  [python]
def test_max_autotune_cutlass_backend_mm_bias(
        self, dynamic: bool = False, max_autotune_gemm_backends: str = "CUTLASS"
    ):
        """
        Make sure autotuning mm in sub processes work without crashes.
        """

        if max_autotune_gemm_backends == "CUTLASS" and torch.version.hip:
            return

        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

        def mm(a, b, bias):
            return torch.nn.functional.linear(a, b, bias)

        a = torch.randn(2048, 4096).cuda().half()
        bias = torch.randn(2048).cuda().half()

        with config.patch(
            {
                "max_autotune": True,
                "autotune_in_subproc": False,
                "max_autotune_gemm_backends": max_autotune_gemm_backends,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 2,
            }
        ):
            Y = mm(a, a, bias)
            Y_compiled = torch.compile(mm, dynamic=dynamic)(a, a, bias)
            torch.testing.assert_close(Y_compiled, Y, atol=1e-1, rtol=1e-1)
//  [python]
def _test4d(self, align_corners):
        H, W = 4, 5
        # Construct a grid to obtain indentity sampling
        input = torch.randn(H * W).view(1, 1, H, W).float()
        coords = torch.meshgrid(torch.arange(H), torch.arange(W))
        coords = torch.stack(coords[::-1], dim=-1).float()[None]
        if not align_corners:
            coords = coords + 0.5
        sampled_input = bilinear_sampler(input, coords, align_corners=align_corners)
        torch.testing.assert_allclose(input, sampled_input)
//  [python]
def _test4d(self, align_corners):
        H, W = 4, 5
        # Construct a grid to obtain indentity sampling
        input = torch.randn(H * W).view(1, 1, H, W).float()
        coords = torch.meshgrid(torch.arange(H), torch.arange(W))
        coords = torch.stack(coords[::-1], dim=-1).float()[None]
        if not align_corners:
            coords = coords + 0.5
        sampled_input = bilinear_sampler(input, coords, align_corners=align_corners)
        torch.testing.assert_close(input, sampled_input)
//  [python]
def test_max_autotune_cutlass_backend_int_mm(
        self, dynamic: bool, max_autotune_gemm_backends: str
    ):
        """
        Make sure autotuning mm in sub processes work without crashes.
        """

        if "CUTLASS" in max_autotune_gemm_backends.upper() and torch.version.hip:
            return

        def mm(a, b):
            return torch._int_mm(a, b)

        # CUTLASS only supports row-major/column-major combination of
        # layouts for this operation, thus the transpose of tensor b
        # (on the other side, Triton at the moment doesn't support
        # this combination, so it's excluded from the test).  Also,
        # for CUTLASS alignment requirements, number of columns in
        # both tensors has to be divisible by 16.
        a = torch.randint(0, 5, (100, 16), dtype=torch.int8).cuda()
        b = torch.randint(0, 5, (32, 16), dtype=torch.int8).cuda().T

        with config.patch(
            {
                "max_autotune": True,
                "autotune_in_subproc": False,
                "max_autotune_gemm_backends": max_autotune_gemm_backends,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 2,
            }
        ):
            Y_compiled = torch.compile(mm, dynamic=dynamic)(a, b)
            Y = mm(a, b)
            t.assert_close(Y_compiled, Y)
//  [python]
def test_max_autotune_cutlass_backend_int_mm(
        self, dynamic: bool, max_autotune_gemm_backends: str
    ):
        """
        Make sure autotuning mm in sub processes work without crashes.
        """

        if "CUTLASS" in max_autotune_gemm_backends.upper() and torch.version.hip:
            return

        def mm(a, b):
            return torch._int_mm(a, b)

        # CUTLASS only supports row-major/column-major combination of
        # layouts for this operation, thus the transpose of tensor b
        # (on the other side, Triton at the moment doesn't support
        # this combination, so it's excluded from the test).  Also,
        # for CUTLASS alignment requirements, number of columns in
        # both tensors has to be divisible by 16.
        a = torch.randint(0, 5, (100, 16), dtype=torch.int8).cuda()
        b = torch.randint(0, 5, (32, 16), dtype=torch.int8).cuda().T

        with config.patch(
            {
                "max_autotune": True,
                "autotune_in_subproc": False,
                "max_autotune_gemm_backends": max_autotune_gemm_backends,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 2,
            }
        ):
            Y_compiled = torch.compile(mm, dynamic=dynamic)(a, b)
            Y = mm(a, b)
            torch.testing.assert_close(Y_compiled, Y)
//  [python]
def test_max_autotune_cutlass_backend_mixed_mm(
        self, dynamic: bool, max_autotune_gemm_backends: str
    ):
        """
        Make sure autotuning mm in sub processes work without crashes.
        """

        if max_autotune_gemm_backends == "CUTLASS" and torch.version.hip:
            return

        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

        def mm(a, b):
            return torch.mm(a, b.to(torch.half))

        # CUTLASS only supports row-major/column-major combination of
        # layouts for this operation, thus the transpose of tensor b.
        # Also, for CUTLASS alignment requirements, number of columns
        # of the first tensor has to be divisible by 16.
        a = torch.randn(100, 16).cuda().half()
        b = torch.randint(0, 5, (100, 16), dtype=torch.int8).cuda().T

        with config.patch(
            {
                "max_autotune": True,
                "autotune_in_subproc": False,
                "max_autotune_gemm_backends": max_autotune_gemm_backends,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 2,
                "use_mixed_mm": True,
            }
        ):
            Y_compiled = torch.compile(mm, dynamic=dynamic)(a, b)
            Y = mm(a, b)
            self.assertEqual(Y_compiled, Y)
//  [python]
def test_max_autotune_cutlass_backend_mixed_mm(
        self, dynamic: bool, max_autotune_gemm_backends: str
    ):
        """
        Make sure autotuning mm in sub processes work without crashes.
        """

        if max_autotune_gemm_backends == "CUTLASS" and torch.version.hip:
            return

        torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False

        def mm(a, b):
            return torch.mm(a, b.to(torch.half))

        # CUTLASS only supports row-major/column-major combination of
        # layouts for this operation, thus the transpose of tensor b.
        # Also, for CUTLASS alignment requirements, number of columns
        # of the first tensor has to be divisible by 16.
        a = torch.randn(100, 16).cuda().half()
        b = torch.randint(0, 5, (100, 16), dtype=torch.int8).cuda().T

        with config.patch(
            {
                "max_autotune": True,
                "autotune_in_subproc": False,
                "max_autotune_gemm_backends": max_autotune_gemm_backends,
                "cuda.cutlass_dir": _CUTLASS_DIR,
                "cuda.cutlass_max_profiling_configs": 2,
                "use_mixed_mm": True,
            }
        ):
            Y_compiled = torch.compile(mm, dynamic=dynamic)(a, b)
            Y = mm(a, b)
            torch.testing.assert_close(Y_compiled, Y)
//  [python]
def assert_expected(
    actual: Any,
    expected: Any,
    rtol: float = 1e-5,
    atol: float = 1e-8,
    check_device: bool = 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: float = 1e-5,
    atol: float = 1e-8,
    check_device: bool = True,
):
    torch.testing.assert_close(
        actual,
        expected,
        rtol=rtol,
        atol=atol,
        check_device=check_device,
        msg=f"actual: {actual}, expected: {expected}",
    )
//  [python]
def _test5d(self, align_corners):
        T, H, W = 3, 4, 5
        # Construct a grid to obtain indentity sampling
        input = torch.randn(H * W).view(1, 1, H, W).float()
        input = torch.stack([input, input + 1, input + 2], dim=2)
        coords = torch.meshgrid(torch.arange(T), torch.arange(W), torch.arange(H))
        coords = torch.stack(coords, dim=-1).float().permute(0, 2, 1, 3)[None]

        if not align_corners:
            coords = coords + 0.5
        sampled_input = bilinear_sampler(input, coords, align_corners=align_corners)
        torch.testing.assert_allclose(input, sampled_input)