jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_linear_cutlass(self, device, dtype):

        def run_test(batch_shape, m, n, k, device, dtype, dtype_out, add_bias, activation, rtol, atol):
            weight = rand_sparse_semi_structured(m, k, dtype, device)
            input = make_tensor((*batch_shape, n, k), dtype=dtype, device=device)
            bias = make_tensor((m,), dtype=dtype_out, device=device) if add_bias else None

            dtype_dense = torch.float32
            input_dense = input.to(dtype_dense)
            weight_dense = weight.to(dtype_dense)
            bias_dense = bias.to(dtype_dense) if add_bias else None
            output0 = torch.nn.functional.linear(input_dense, weight_dense, bias=bias_dense)
            if activation == "relu":
                relu = torch.nn.ReLU()
                output0 = relu(output0)
            elif activation == "silu":
                silu = torch.nn.SiLU()
                output0 = silu(output0)

            compressed = to_sparse_semi_structured(weight)

            weight_sparse = compressed.values()
            meta = compressed.indices()

            output1 = torch._sparse_semi_structured_linear(input, weight_sparse, meta, bias=bias, activation=activation,
                                                           out_dtype=dtype_out if dtype == torch.int8 else None)
            torch.testing.assert_allclose(output1.to(dtype_dense), output0, rtol=rtol, atol=atol)

        if dtype == torch.float32:
            # Inputs are converted to TF32 internally for sparse GEMM,
            # so make dense GEMM to do the same for matching results.
            orig = torch.backends.cuda.matmul.allow_tf32
            torch.backends.cuda.matmul.allow_tf32 = True

        batch_shapes = [[], [3], [3, 1]]
        dtype_out = {torch.int8: torch.int32, torch.half: torch.half, torch.bfloat16: torch.bfloat16, torch.float32: torch.float32}
        activations = [None, "relu", "silu"]
        rtol, atol = 1e-3, 1e-3
        if dtype == torch.bfloat16:
            rtol, atol = 5e-3, 5e-3
        elif dtype == torch.float32:
            rtol, atol = 1e-3, 75e-2
        for batch_shape, m, n, k, add_bias, activation in \
                itertools.product(batch_shapes, range(3), range(3), range(3), (False, True), activations):
            if activation == "silu" and dtype == torch.int8:
                continue  # SiLU not supported for integer inputs

            m = 2 ** m * 32
            n = 2 ** n * 32
            k = 2 ** k * 128
            run_test(batch_shape, m, n, k, device, dtype, dtype_out[dtype], add_bias, activation, rtol, atol)

        if dtype == torch.float32:
            torch.backends.cuda.matmul.allow_tf32 = orig
//  [python]
def test_linear_cutlass(self, device, dtype):

        def run_test(batch_shape, m, n, k, device, dtype, dtype_out, add_bias, activation, rtol, atol):
            weight = rand_sparse_semi_structured(m, k, dtype, device)
            input = make_tensor((*batch_shape, n, k), dtype=dtype, device=device)
            bias = make_tensor((m,), dtype=dtype_out, device=device) if add_bias else None

            dtype_dense = torch.float32
            input_dense = input.to(dtype_dense)
            weight_dense = weight.to(dtype_dense)
            bias_dense = bias.to(dtype_dense) if add_bias else None
            output0 = torch.nn.functional.linear(input_dense, weight_dense, bias=bias_dense)
            if activation == "relu":
                relu = torch.nn.ReLU()
                output0 = relu(output0)
            elif activation == "silu":
                silu = torch.nn.SiLU()
                output0 = silu(output0)

            compressed = to_sparse_semi_structured(weight)

            weight_sparse = compressed.values()
            meta = compressed.indices()

            output1 = torch._sparse_semi_structured_linear(input, weight_sparse, meta, bias=bias, activation=activation,
                                                           out_dtype=dtype_out if dtype == torch.int8 else None)
            torch.testing.assert_close(output1.to(dtype_dense), output0, rtol=rtol, atol=atol)

        if dtype == torch.float32:
            # Inputs are converted to TF32 internally for sparse GEMM,
            # so make dense GEMM to do the same for matching results.
            orig = torch.backends.cuda.matmul.allow_tf32
            torch.backends.cuda.matmul.allow_tf32 = True

        batch_shapes = [[], [3], [3, 1]]
        dtype_out = {torch.int8: torch.int32, torch.half: torch.half, torch.bfloat16: torch.bfloat16, torch.float32: torch.float32}
        activations = [None, "relu", "silu"]
        rtol, atol = 1e-3, 1e-3
        if dtype == torch.bfloat16:
            rtol, atol = 5e-3, 5e-3
        elif dtype == torch.float32:
            rtol, atol = 1e-3, 75e-2
        for batch_shape, m, n, k, add_bias, activation in \
                itertools.product(batch_shapes, range(3), range(3), range(3), (False, True), activations):
            if activation == "silu" and dtype == torch.int8:
                continue  # SiLU not supported for integer inputs

            m = 2 ** m * 32
            n = 2 ** n * 32
            k = 2 ** k * 128
            run_test(batch_shape, m, n, k, device, dtype, dtype_out[dtype], add_bias, activation, rtol, atol)

        if dtype == torch.float32:
            torch.backends.cuda.matmul.allow_tf32 = orig
//  [python]
def test_token_accuracy(ignore_index: int, vocab_size: int):
    batch_size = int(1e6)
    torchmetrics_token_acc = TokenAccuracy(ignore_index=ignore_index)
    generated_preds = torch.rand((batch_size, vocab_size))
    true_labels = torch.randint(low=0, high=vocab_size - 1, size=(batch_size,))

    # Randomly insert ignore_index into the labels
    labels_mask = torch.rand((batch_size,))
    labels_mask[labels_mask > 0.8] = 1
    labels_mask[labels_mask <= 0.8] = 0
    labels_mask = labels_mask.bool()
    true_labels[labels_mask] = ignore_index

    true_labels = true_labels.float()
    generated_preds = generated_preds.float()

    torchmetrics_token_acc.update(generated_preds, true_labels)
    final_acc = torchmetrics_token_acc.compute()

    expected_random_acc_tensor = torch.tensor(1.0 / vocab_size)
    t.assert_close(final_acc, expected_random_acc_tensor)
//  [python]
def test_token_accuracy(ignore_index: int, vocab_size: int):
    batch_size = int(1e6)
    torchmetrics_token_acc = TokenAccuracy(ignore_index=ignore_index)
    generated_preds = torch.rand((batch_size, vocab_size))
    true_labels = torch.randint(low=0, high=vocab_size - 1, size=(batch_size,))

    # Randomly insert ignore_index into the labels
    labels_mask = torch.rand((batch_size,))
    labels_mask[labels_mask > 0.8] = 1
    labels_mask[labels_mask <= 0.8] = 0
    labels_mask = labels_mask.bool()
    true_labels[labels_mask] = ignore_index

    true_labels = true_labels.float()
    generated_preds = generated_preds.float()

    torchmetrics_token_acc.update(generated_preds, true_labels)
    final_acc = torchmetrics_token_acc.compute()

    expected_random_acc_tensor = torch.tensor(1.0 / vocab_size)
    torch.testing.assert_close(final_acc, expected_random_acc_tensor)
//  [python]
def test_out(make_args, out, nnc=nnc_add, aten=torch.add):
    def out_setup(n):
        args = make_args(n)
        result_aten = out(n)
        result_nnc = out(n)
        aten(*args, out=result_aten)
        nnc(*args, out=result_nnc)
        t.assert_close(result_aten, result_nnc)
        result = out(n)
        return (lambda: nnc(*args, out=result), lambda: aten(*args, out=result))

    return benchmark_loop(out_setup)
//  [python]
def test_out(make_args, out, nnc=nnc_add, aten=torch.add):
    def out_setup(n):
        args = make_args(n)
        result_aten = out(n)
        result_nnc = out(n)
        aten(*args, out=result_aten)
        nnc(*args, out=result_nnc)
        torch.testing.assert_close(result_aten, result_nnc)
        result = out(n)
        return (lambda: nnc(*args, out=result), lambda: aten(*args, out=result))

    return benchmark_loop(out_setup)
//  [python]
def test_backwards(make_args, nnc=nnc_add, aten=torch.add):
    def backwards_setup(n):
        args = make_args(n)
        (grad_var,) = (a for a in args if a.requires_grad)
        aten(*args).sum().backward()
        correct = grad_var.grad.clone()
        grad_var.grad.zero_()
        nnc(*args).sum().backward()
        torch.testing.assert_allclose(correct, grad_var.grad)
        return (
            lambda: nnc(*args).sum().backward(),
            lambda: aten(*args).sum().backward(),
        )

    return benchmark_loop(backwards_setup)
//  [python]
def test_backwards(make_args, nnc=nnc_add, aten=torch.add):
    def backwards_setup(n):
        args = make_args(n)
        (grad_var,) = (a for a in args if a.requires_grad)
        aten(*args).sum().backward()
        correct = grad_var.grad.clone()
        grad_var.grad.zero_()
        nnc(*args).sum().backward()
        torch.testing.assert_close(correct, grad_var.grad)
        return (
            lambda: nnc(*args).sum().backward(),
            lambda: aten(*args).sum().backward(),
        )

    return benchmark_loop(backwards_setup)
//  [python]
def test_inplace(make_args, nnc=nnc_add, aten=torch.add):
    def inplace_setup(n):
        a, b = make_args(n)
        result_aten = torch.clone(a)
        result_nnc = torch.clone(a)
        nnc(result_nnc, b, out=result_nnc)
        aten(result_aten, b, out=result_aten)
        torch.testing.assert_allclose(result_aten, result_nnc)
        return (lambda: nnc(a, b, out=a), lambda: aten(a, b, out=a))

    return benchmark_loop(inplace_setup)
//  [python]
def test_inplace(make_args, nnc=nnc_add, aten=torch.add):
    def inplace_setup(n):
        a, b = make_args(n)
        result_aten = torch.clone(a)
        result_nnc = torch.clone(a)
        nnc(result_nnc, b, out=result_nnc)
        aten(result_aten, b, out=result_aten)
        torch.testing.assert_close(result_aten, result_nnc)
        return (lambda: nnc(a, b, out=a), lambda: aten(a, b, out=a))

    return benchmark_loop(inplace_setup)
//  [python]
def test_rvq_eval(self, num_codebooks: int):
        """Simple test to confirm that the RVQ module can be instantiated and run,
        and that forward produces the same result as encode-decode.
        """
        # instantiate and set in eval mode
        rvq = ResidualVectorQuantizer(num_codebooks=num_codebooks, codebook_dim=self.codebook_dim)
        rvq.eval()

        for i in range(self.num_examples):
            inputs = torch.randn([self.batch_size, self.codebook_dim, self.max_len])
            input_len = torch.tensor([self.max_len] * self.batch_size, dtype=torch.int32)

            # apply forward
            dequantized_fw, indices_fw, commit_loss = rvq(inputs=inputs, input_len=input_len)

            # make sure the commit loss is zero
            assert commit_loss == 0.0, f'example {i}: commit_loss is {commit_loss}, expected 0.0'

            # encode-decode
            indices_enc = rvq.encode(inputs=inputs, input_len=input_len)
            dequantized_dec = rvq.decode(indices=indices_enc, input_len=input_len)

            # make sure the results are the same
            t.assert_allclose(indices_enc, indices_fw, msg=f'example {i}: indices mismatch')
            t.assert_allclose(dequantized_dec, dequantized_fw, msg=f'example {i}: dequantized mismatch')
//  [python]
def test_rvq_eval(self, num_codebooks: int):
        """Simple test to confirm that the RVQ module can be instantiated and run,
        and that forward produces the same result as encode-decode.
        """
        # instantiate and set in eval mode
        rvq = ResidualVectorQuantizer(num_codebooks=num_codebooks, codebook_dim=self.codebook_dim)
        rvq.eval()

        for i in range(self.num_examples):
            inputs = torch.randn([self.batch_size, self.codebook_dim, self.max_len])
            input_len = torch.tensor([self.max_len] * self.batch_size, dtype=torch.int32)

            # apply forward
            dequantized_fw, indices_fw, commit_loss = rvq(inputs=inputs, input_len=input_len)

            # make sure the commit loss is zero
            assert commit_loss == 0.0, f'example {i}: commit_loss is {commit_loss}, expected 0.0'

            # encode-decode
            indices_enc = rvq.encode(inputs=inputs, input_len=input_len)
            dequantized_dec = rvq.decode(indices=indices_enc, input_len=input_len)

            # make sure the results are the same
            torch.testing.assert_close(indices_enc, indices_fw, msg=f'example {i}: indices mismatch')
            torch.testing.assert_close(dequantized_dec, dequantized_fw, msg=f'example {i}: dequantized mismatch')
//  [python]
def test_fsq_eval(self, num_levels: list):
        """Simple test to confirm that the FSQ module can be instantiated and run,
        and that forward produces the same result as encode-decode.
        """
        fsq = FiniteScalarQuantizer(num_levels=num_levels)

        for i in range(self.num_examples):
            inputs = torch.randn([self.batch_size, fsq.codebook_dim, self.max_len])
            input_len = torch.tensor([self.max_len] * self.batch_size, dtype=torch.int32)

            # apply forward
            dequantized_fw, indices_fw = fsq(inputs=inputs, input_len=input_len)

            assert dequantized_fw.max() <= 1.0, f'example {i}: dequantized_fw.max() is {dequantized_fw.max()}'
            assert dequantized_fw.min() >= -1.0, f'example {i}: dequantized_fw.min() is {dequantized_fw.min()}'

            # encode-decode
            indices_enc = fsq.encode(inputs=inputs, input_len=input_len)
            dequantized_dec = fsq.decode(indices=indices_enc, input_len=input_len)

            # make sure the results are the same
            torch.testing.assert_allclose(indices_enc, indices_fw, msg=f'example {i}: indices mismatch')
            torch.testing.assert_allclose(dequantized_dec, dequantized_fw, msg=f'example {i}: dequantized mismatch')
//  [python]
def test_fsq_eval(self, num_levels: list):
        """Simple test to confirm that the FSQ module can be instantiated and run,
        and that forward produces the same result as encode-decode.
        """
        fsq = FiniteScalarQuantizer(num_levels=num_levels)

        for i in range(self.num_examples):
            inputs = torch.randn([self.batch_size, fsq.codebook_dim, self.max_len])
            input_len = torch.tensor([self.max_len] * self.batch_size, dtype=torch.int32)

            # apply forward
            dequantized_fw, indices_fw = fsq(inputs=inputs, input_len=input_len)

            assert dequantized_fw.max() <= 1.0, f'example {i}: dequantized_fw.max() is {dequantized_fw.max()}'
            assert dequantized_fw.min() >= -1.0, f'example {i}: dequantized_fw.min() is {dequantized_fw.min()}'

            # encode-decode
            indices_enc = fsq.encode(inputs=inputs, input_len=input_len)
            dequantized_dec = fsq.decode(indices=indices_enc, input_len=input_len)

            # make sure the results are the same
            torch.testing.assert_close(indices_enc, indices_fw, msg=f'example {i}: indices mismatch')
            torch.testing.assert_close(dequantized_dec, dequantized_fw, msg=f'example {i}: dequantized mismatch')
//  [python]
def _test_fuse_split_cat_rearrange(self, M, N, split, remove_split=True):
        dtype = "float16"
        M = IntImm(M)
        N = IntImm(N)

        input_1 = Tensor(
            shape=[M, N],
            name="input_1",
            is_input=True,
        )
        split_2 = ops.split()(input_1, split, 0)
        concatenate_3 = ops.concatenate()(split_2[::-1], 0)

        # Set outputs
        concatenate_3._attrs["name"] = "output_0"
        concatenate_3._attrs["is_output"] = True
        # Compile
        model = compile_model(
            concatenate_3, detect_target(), "./tmp", self._testMethodName
        )
        # Check that split was removed
        self.assertEqual(
            graph_has_op(model.debug_sorted_graph, "split"), not remove_split
        )
        # Run
        input_1 = get_random_torch_tensor((M.value(), N.value()), dtype=dtype)
        # Compare
        split_pt = torch.split(input_1, split, 0)
        y_pt = torch.cat(
            [split_pt[1], split_pt[0]],
            0,
        )
        y_ait = torch.empty_like(y_pt)
        model.run_with_tensors(
            {"input_1": input_1},
            [y_ait],
        )
        torch.testing.assert_allclose(y_ait, y_pt, atol=0, rtol=0)
//  [python]
def _test_fuse_split_cat_rearrange(self, M, N, split, remove_split=True):
        dtype = "float16"
        M = IntImm(M)
        N = IntImm(N)

        input_1 = Tensor(
            shape=[M, N],
            name="input_1",
            is_input=True,
        )
        split_2 = ops.split()(input_1, split, 0)
        concatenate_3 = ops.concatenate()(split_2[::-1], 0)

        # Set outputs
        concatenate_3._attrs["name"] = "output_0"
        concatenate_3._attrs["is_output"] = True
        # Compile
        model = compile_model(
            concatenate_3, detect_target(), "./tmp", self._testMethodName
        )
        # Check that split was removed
        self.assertEqual(
            graph_has_op(model.debug_sorted_graph, "split"), not remove_split
        )
        # Run
        input_1 = get_random_torch_tensor((M.value(), N.value()), dtype=dtype)
        # Compare
        split_pt = torch.split(input_1, split, 0)
        y_pt = torch.cat(
            [split_pt[1], split_pt[0]],
            0,
        )
        y_ait = torch.empty_like(y_pt)
        model.run_with_tensors(
            {"input_1": input_1},
            [y_ait],
        )
        torch.testing.assert_close(y_ait, y_pt, atol=0, rtol=0)
//  [python]
def test_fuse_split_cat_dim1(self):
        dtype = "float16"
        M = IntImm(512)
        N = IntImm(512)

        input_1 = Tensor(
            shape=[M, N],
            name="input_1",
            is_input=True,
        )
        split_2 = ops.split()(input_1, int(N.value() / 2), 1)
        concatenate_3 = ops.concatenate()(split_2[::-1], 1)

        # Set outputs
        concatenate_3._attrs["name"] = "output_0"
        concatenate_3._attrs["is_output"] = True
        # Compile
        model = compile_model(
            concatenate_3, detect_target(), "./tmp", self._testMethodName
        )
        # Check that split was removed
        self.assertFalse(graph_has_op(model.debug_sorted_graph, "split"))
        # Run
        input_1 = get_random_torch_tensor((M.value(), N.value()), dtype=dtype)
        # Compare
        split_pt = torch.split(input_1, int(N.value() / 2), 1)
        y_pt = torch.cat(
            split_pt[::-1],
            1,
        )
        y_ait = torch.empty_like(y_pt)
        model.run_with_tensors(
            {"input_1": input_1},
            [y_ait],
        )
        torch.testing.assert_allclose(y_ait, y_pt, atol=0, rtol=0)
//  [python]
def test_fuse_split_cat_dim1(self):
        dtype = "float16"
        M = IntImm(512)
        N = IntImm(512)

        input_1 = Tensor(
            shape=[M, N],
            name="input_1",
            is_input=True,
        )
        split_2 = ops.split()(input_1, int(N.value() / 2), 1)
        concatenate_3 = ops.concatenate()(split_2[::-1], 1)

        # Set outputs
        concatenate_3._attrs["name"] = "output_0"
        concatenate_3._attrs["is_output"] = True
        # Compile
        model = compile_model(
            concatenate_3, detect_target(), "./tmp", self._testMethodName
        )
        # Check that split was removed
        self.assertFalse(graph_has_op(model.debug_sorted_graph, "split"))
        # Run
        input_1 = get_random_torch_tensor((M.value(), N.value()), dtype=dtype)
        # Compare
        split_pt = torch.split(input_1, int(N.value() / 2), 1)
        y_pt = torch.cat(
            split_pt[::-1],
            1,
        )
        y_ait = torch.empty_like(y_pt)
        model.run_with_tensors(
            {"input_1": input_1},
            [y_ait],
        )
        torch.testing.assert_close(y_ait, y_pt, atol=0, rtol=0)
//  [python]
def test_fuse_split_cat_different_dims(self):
        """Splitting and then concatting on different dims is not
        expected to be optimized currently."""
        dtype = "float16"
        M = IntImm(512)
        N = IntImm(512)

        input_1 = Tensor(
            shape=[M, N],
            name="input_1",
            is_input=True,
        )
        split_2 = ops.split()(input_1, int(M.value() / 2), 0)
        concatenate_3 = ops.concatenate()(split_2[::-1], 1)

        # Set outputs
        concatenate_3._attrs["name"] = "output_0"
        concatenate_3._attrs["is_output"] = True
        # Compile
        model = compile_model(
            concatenate_3, detect_target(), "./tmp", self._testMethodName
        )
        # Check that split was not removed because the dims are different
        self.assertTrue(graph_has_op(model.debug_sorted_graph, "split"))
        # Run
        input_1 = get_random_torch_tensor((M.value(), N.value()), dtype=dtype)
        # Compare
        split_pt = torch.split(input_1, int(M.value() / 2), 0)
        y_pt = torch.cat(
            split_pt[::-1],
            1,
        )
        y_ait = torch.empty_like(y_pt)
        model.run_with_tensors(
            {"input_1": input_1},
            [y_ait],
        )
        torch.testing.assert_allclose(y_ait, y_pt, atol=0, rtol=0)
//  [python]
def test_fuse_split_cat_different_dims(self):
        """Splitting and then concatting on different dims is not
        expected to be optimized currently."""
        dtype = "float16"
        M = IntImm(512)
        N = IntImm(512)

        input_1 = Tensor(
            shape=[M, N],
            name="input_1",
            is_input=True,
        )
        split_2 = ops.split()(input_1, int(M.value() / 2), 0)
        concatenate_3 = ops.concatenate()(split_2[::-1], 1)

        # Set outputs
        concatenate_3._attrs["name"] = "output_0"
        concatenate_3._attrs["is_output"] = True
        # Compile
        model = compile_model(
            concatenate_3, detect_target(), "./tmp", self._testMethodName
        )
        # Check that split was not removed because the dims are different
        self.assertTrue(graph_has_op(model.debug_sorted_graph, "split"))
        # Run
        input_1 = get_random_torch_tensor((M.value(), N.value()), dtype=dtype)
        # Compare
        split_pt = torch.split(input_1, int(M.value() / 2), 0)
        y_pt = torch.cat(
            split_pt[::-1],
            1,
        )
        y_ait = torch.empty_like(y_pt)
        model.run_with_tensors(
            {"input_1": input_1},
            [y_ait],
        )
        torch.testing.assert_close(y_ait, y_pt, atol=0, rtol=0)
//  [python]
def test_fuse_split_cat_reuse(self):
        """Use a split output twice in the concatenate op."""
        dtype = "float16"
        M = IntImm(512)
        N = IntImm(512)

        input_1 = Tensor(
            shape=[M, N],
            name="input_1",
            is_input=True,
        )
        split_2 = ops.split()(input_1, int(M.value() / 2), 0)
        concatenate_3 = ops.concatenate()([split_2[1], split_2[0], split_2[1]], 0)

        # Set outputs
        concatenate_3._attrs["name"] = "output_0"
        concatenate_3._attrs["is_output"] = True
        # Compile
        model = compile_model(
            concatenate_3, detect_target(), "./tmp", self._testMethodName
        )
        # Check that split was removed
        self.assertFalse(graph_has_op(model.debug_sorted_graph, "split"))
        # Run
        input_1 = get_random_torch_tensor((M.value(), N.value()), dtype=dtype)
        # Compare
        split_pt = torch.split(input_1, int(M.value() / 2), 0)
        y_pt = torch.cat(
            [split_pt[1], split_pt[0], split_pt[1]],
            0,
        )
        y_ait = torch.empty_like(y_pt)
        model.run_with_tensors(
            {"input_1": input_1},
            [y_ait],
        )
        t.assert_allclose(y_ait, y_pt, atol=0, rtol=0)
//  [python]
def test_fuse_split_cat_reuse(self):
        """Use a split output twice in the concatenate op."""
        dtype = "float16"
        M = IntImm(512)
        N = IntImm(512)

        input_1 = Tensor(
            shape=[M, N],
            name="input_1",
            is_input=True,
        )
        split_2 = ops.split()(input_1, int(M.value() / 2), 0)
        concatenate_3 = ops.concatenate()([split_2[1], split_2[0], split_2[1]], 0)

        # Set outputs
        concatenate_3._attrs["name"] = "output_0"
        concatenate_3._attrs["is_output"] = True
        # Compile
        model = compile_model(
            concatenate_3, detect_target(), "./tmp", self._testMethodName
        )
        # Check that split was removed
        self.assertFalse(graph_has_op(model.debug_sorted_graph, "split"))
        # Run
        input_1 = get_random_torch_tensor((M.value(), N.value()), dtype=dtype)
        # Compare
        split_pt = torch.split(input_1, int(M.value() / 2), 0)
        y_pt = torch.cat(
            [split_pt[1], split_pt[0], split_pt[1]],
            0,
        )
        y_ait = torch.empty_like(y_pt)
        model.run_with_tensors(
            {"input_1": input_1},
            [y_ait],
        )
        torch.testing.assert_close(y_ait, y_pt, atol=0, rtol=0)
//  [python]
def test_group_fsq_eval(self, num_groups: int, num_levels_per_group: int):
        """Simple test to confirm that the group FSQ module can be instantiated and run,
        and that forward produces the same result as encode-decode.
        """
        # Test inference with group FSQ
        # instantiate
        gfsq = GroupFiniteScalarQuantizer(num_groups=num_groups, num_levels_per_group=num_levels_per_group)

        for i in range(self.num_examples):
            inputs = torch.randn([self.batch_size, gfsq.codebook_dim, self.max_len])
            input_len = torch.tensor([self.max_len] * self.batch_size, dtype=torch.int32)

            # apply forward
            dequantized_fw, indices_fw = gfsq(inputs=inputs, input_len=input_len)

            # encode-decode
            indices_enc = gfsq.encode(inputs=inputs, input_len=input_len)
            dequantized_dec = gfsq.decode(indices=indices_enc, input_len=input_len)

            # make sure the results are the same
            t.assert_close(indices_enc, indices_fw, msg=f'example {i}: indices mismatch')
            t.assert_close(dequantized_dec, dequantized_fw, msg=f'example {i}: dequantized mismatch')

            # apply individual FSQs and make sure the results are the same
            inputs_grouped = inputs.chunk(num_groups, dim=1)
            dequantized_fw_grouped = dequantized_fw.chunk(num_groups, dim=1)
            indices_fw_grouped = indices_fw.chunk(num_groups, dim=0)

            for g in range(num_groups):
                dequantized, indices = gfsq.fsqs[g](inputs=inputs_grouped[g], input_len=input_len)
                t.assert_close(
                    dequantized, dequantized_fw_grouped[g], msg=f'example {i}: dequantized mismatch for group {g}'
                )
                t.assert_close(
                    indices, indices_fw_grouped[g], msg=f'example {i}: indices mismatch for group {g}'
                )
//  [python]
def test_group_fsq_eval(self, num_groups: int, num_levels_per_group: int):
        """Simple test to confirm that the group FSQ module can be instantiated and run,
        and that forward produces the same result as encode-decode.
        """
        # Test inference with group FSQ
        # instantiate
        gfsq = GroupFiniteScalarQuantizer(num_groups=num_groups, num_levels_per_group=num_levels_per_group)

        for i in range(self.num_examples):
            inputs = torch.randn([self.batch_size, gfsq.codebook_dim, self.max_len])
            input_len = torch.tensor([self.max_len] * self.batch_size, dtype=torch.int32)

            # apply forward
            dequantized_fw, indices_fw = gfsq(inputs=inputs, input_len=input_len)

            # encode-decode
            indices_enc = gfsq.encode(inputs=inputs, input_len=input_len)
            dequantized_dec = gfsq.decode(indices=indices_enc, input_len=input_len)

            # make sure the results are the same
            torch.testing.assert_close(indices_enc, indices_fw, msg=f'example {i}: indices mismatch')
            torch.testing.assert_close(dequantized_dec, dequantized_fw, msg=f'example {i}: dequantized mismatch')

            # apply individual FSQs and make sure the results are the same
            inputs_grouped = inputs.chunk(num_groups, dim=1)
            dequantized_fw_grouped = dequantized_fw.chunk(num_groups, dim=1)
            indices_fw_grouped = indices_fw.chunk(num_groups, dim=0)

            for g in range(num_groups):
                dequantized, indices = gfsq.fsqs[g](inputs=inputs_grouped[g], input_len=input_len)
                torch.testing.assert_close(
                    dequantized, dequantized_fw_grouped[g], msg=f'example {i}: dequantized mismatch for group {g}'
                )
                torch.testing.assert_close(
                    indices, indices_fw_grouped[g], msg=f'example {i}: indices mismatch for group {g}'
                )
//  [python]
def test_group_rvq_eval(self, num_groups: int, num_codebooks: int):
        """Simple test to confirm that the group RVQ module can be instantiated and run,
        and that forward produces the same result as encode-decode.
        """
        if num_groups > num_codebooks:
            # Expected to fail if num_groups is lager than the total number of codebooks
            with pytest.raises(ValueError):
                _ = GroupResidualVectorQuantizer(
                    num_codebooks=num_codebooks, num_groups=num_groups, codebook_dim=self.codebook_dim
                )
        else:
            # Test inference with group RVQ
            # instantiate and set in eval mode
            grvq = GroupResidualVectorQuantizer(
                num_codebooks=num_codebooks, num_groups=num_groups, codebook_dim=self.codebook_dim
            )
            grvq.eval()

            for i in range(self.num_examples):
                inputs = torch.randn([self.batch_size, self.codebook_dim, self.max_len])
                input_len = torch.tensor([self.max_len] * self.batch_size, dtype=torch.int32)

                # apply forward
                dequantized_fw, indices_fw, commit_loss = grvq(inputs=inputs, input_len=input_len)

                # make sure the commit loss is zero
                assert commit_loss == 0.0, f'example {i}: commit_loss is {commit_loss}, expected 0.0'

                # encode-decode
                indices_enc = grvq.encode(inputs=inputs, input_len=input_len)
                dequantized_dec = grvq.decode(indices=indices_enc, input_len=input_len)

                # make sure the results are the same
                t.assert_close(indices_enc, indices_fw, msg=f'example {i}: indices mismatch')
                t.assert_close(dequantized_dec, dequantized_fw, msg=f'example {i}: dequantized mismatch')

                # apply individual RVQs and make sure the results are the same
                inputs_grouped = inputs.chunk(num_groups, dim=1)
                dequantized_fw_grouped = dequantized_fw.chunk(num_groups, dim=1)
                indices_fw_grouped = indices_fw.chunk(num_groups, dim=0)

                for g in range(num_groups):
                    dequantized, indices, _ = grvq.rvqs[g](inputs=inputs_grouped[g], input_len=input_len)
                    t.assert_close(
                        dequantized, dequantized_fw_grouped[g], msg=f'example {i}: dequantized mismatch for group {g}'
                    )
                    t.assert_close(
                        indices, indices_fw_grouped[g], msg=f'example {i}: indices mismatch for group {g}'
                    )
//  [python]
def test_group_rvq_eval(self, num_groups: int, num_codebooks: int):
        """Simple test to confirm that the group RVQ module can be instantiated and run,
        and that forward produces the same result as encode-decode.
        """
        if num_groups > num_codebooks:
            # Expected to fail if num_groups is lager than the total number of codebooks
            with pytest.raises(ValueError):
                _ = GroupResidualVectorQuantizer(
                    num_codebooks=num_codebooks, num_groups=num_groups, codebook_dim=self.codebook_dim
                )
        else:
            # Test inference with group RVQ
            # instantiate and set in eval mode
            grvq = GroupResidualVectorQuantizer(
                num_codebooks=num_codebooks, num_groups=num_groups, codebook_dim=self.codebook_dim
            )
            grvq.eval()

            for i in range(self.num_examples):
                inputs = torch.randn([self.batch_size, self.codebook_dim, self.max_len])
                input_len = torch.tensor([self.max_len] * self.batch_size, dtype=torch.int32)

                # apply forward
                dequantized_fw, indices_fw, commit_loss = grvq(inputs=inputs, input_len=input_len)

                # make sure the commit loss is zero
                assert commit_loss == 0.0, f'example {i}: commit_loss is {commit_loss}, expected 0.0'

                # encode-decode
                indices_enc = grvq.encode(inputs=inputs, input_len=input_len)
                dequantized_dec = grvq.decode(indices=indices_enc, input_len=input_len)

                # make sure the results are the same
                torch.testing.assert_close(indices_enc, indices_fw, msg=f'example {i}: indices mismatch')
                torch.testing.assert_close(dequantized_dec, dequantized_fw, msg=f'example {i}: dequantized mismatch')

                # apply individual RVQs and make sure the results are the same
                inputs_grouped = inputs.chunk(num_groups, dim=1)
                dequantized_fw_grouped = dequantized_fw.chunk(num_groups, dim=1)
                indices_fw_grouped = indices_fw.chunk(num_groups, dim=0)

                for g in range(num_groups):
                    dequantized, indices, _ = grvq.rvqs[g](inputs=inputs_grouped[g], input_len=input_len)
                    torch.testing.assert_close(
                        dequantized, dequantized_fw_grouped[g], msg=f'example {i}: dequantized mismatch for group {g}'
                    )
                    torch.testing.assert_close(
                        indices, indices_fw_grouped[g], msg=f'example {i}: indices mismatch for group {g}'
                    )
//  [python]
def verify_jacobian(self, J1, J2):
        for j1, j2 in zip(J1, J2):
            assert not torch.any(torch.isnan(j1))
            torch.testing.assert_allclose(j1, j2)
//  [python]
def verify_jacobian(self, J1, J2):
        for j1, j2 in zip(J1, J2):
            assert not torch.any(torch.isnan(j1))
            torch.testing.assert_close(j1, j2)
//  [python]
def test_scalar_window(self) -> None:
        """
        Test windowing of scalar metric gives average of previously reported values.
        """
        metric_to_log = torch.tensor([0.1])

        self.scalar.update(
            labels={DefaultTaskInfo.name: metric_to_log},
            predictions={DefaultTaskInfo.name: metric_to_log},
            weights={DefaultTaskInfo.name: metric_to_log},
        )

        metric_to_log = torch.tensor([0.9])

        self.scalar.update(
            labels={DefaultTaskInfo.name: metric_to_log},
            predictions={DefaultTaskInfo.name: metric_to_log},
            weights={DefaultTaskInfo.name: metric_to_log},
        )

        metric = self.scalar.compute()

        actual_window_metric = metric[f"scalar-{DefaultTaskInfo.name}|window_scalar"]

        expected_window_metric = torch.tensor([0.5])

        torch.testing.assert_allclose(
            actual_window_metric,
            expected_window_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_window_metric}, Expected: {expected_window_metric}",
        )
//  [python]
def test_scalar_window(self) -> None:
        """
        Test windowing of scalar metric gives average of previously reported values.
        """
        metric_to_log = torch.tensor([0.1])

        self.scalar.update(
            labels={DefaultTaskInfo.name: metric_to_log},
            predictions={DefaultTaskInfo.name: metric_to_log},
            weights={DefaultTaskInfo.name: metric_to_log},
        )

        metric_to_log = torch.tensor([0.9])

        self.scalar.update(
            labels={DefaultTaskInfo.name: metric_to_log},
            predictions={DefaultTaskInfo.name: metric_to_log},
            weights={DefaultTaskInfo.name: metric_to_log},
        )

        metric = self.scalar.compute()

        actual_window_metric = metric[f"scalar-{DefaultTaskInfo.name}|window_scalar"]

        expected_window_metric = torch.tensor([0.5])

        torch.testing.assert_close(
            actual_window_metric,
            expected_window_metric,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_window_metric}, Expected: {expected_window_metric}",
        )
//  [python]
def test_tensor_to_float_conversion(tmp_path):
    """Check tensor gets converted to float."""

    class TestModel(BoringModel):
        def training_step(self, batch, batch_idx):
            self.log("a", torch.tensor(0.123), prog_bar=True, on_epoch=False)
            self.log("b", torch.tensor([1]), prog_bar=True, on_epoch=False)
            self.log("c", 2, prog_bar=True, on_epoch=False)
            return super().training_step(batch, batch_idx)

    trainer = Trainer(
        default_root_dir=tmp_path, max_epochs=1, limit_train_batches=2, logger=False, enable_checkpointing=False
    )
    trainer.fit(TestModel())

    t.assert_allclose(trainer.progress_bar_metrics["a"], 0.123)
    assert trainer.progress_bar_metrics["b"] == 1.0
    assert trainer.progress_bar_metrics["c"] == 2.0
    pbar = trainer.progress_bar_callback.train_progress_bar
    actual = str(pbar.postfix)
    assert actual.endswith("a=0.123, b=1.000, c=2.000"), actual
//  [python]
def test_tensor_to_float_conversion(tmp_path):
    """Check tensor gets converted to float."""

    class TestModel(BoringModel):
        def training_step(self, batch, batch_idx):
            self.log("a", torch.tensor(0.123), prog_bar=True, on_epoch=False)
            self.log("b", torch.tensor([1]), prog_bar=True, on_epoch=False)
            self.log("c", 2, prog_bar=True, on_epoch=False)
            return super().training_step(batch, batch_idx)

    trainer = Trainer(
        default_root_dir=tmp_path, max_epochs=1, limit_train_batches=2, logger=False, enable_checkpointing=False
    )
    trainer.fit(TestModel())

    torch.testing.assert_close(trainer.progress_bar_metrics["a"], 0.123)
    assert trainer.progress_bar_metrics["b"] == 1.0
    assert trainer.progress_bar_metrics["c"] == 2.0
    pbar = trainer.progress_bar_callback.train_progress_bar
    actual = str(pbar.postfix)
    assert actual.endswith("a=0.123, b=1.000, c=2.000"), actual
//  [python]
def test_scalar(self) -> None:
        """
        Test scalar metric passes through each tensor as is
        """
        metric_to_log = torch.tensor([0.1])

        self.scalar.update(
            labels={DefaultTaskInfo.name: metric_to_log},
            predictions={DefaultTaskInfo.name: metric_to_log},
            weights={DefaultTaskInfo.name: metric_to_log},
        )
        metric = self.scalar.compute()
        actual_metric = metric[f"scalar-{DefaultTaskInfo.name}|lifetime_scalar"]

        torch.testing.assert_allclose(
            actual_metric,
            metric_to_log,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {metric_to_log}",
        )

        # Pass through second tensor with different value
        # check we get the value back with no averaging or any differences

        metric_to_log = torch.tensor([0.9])

        self.scalar.update(
            labels={DefaultTaskInfo.name: metric_to_log},
            predictions={DefaultTaskInfo.name: metric_to_log},
            weights={DefaultTaskInfo.name: metric_to_log},
        )
        metric = self.scalar.compute()
        actual_metric = metric[f"scalar-{DefaultTaskInfo.name}|lifetime_scalar"]

        torch.testing.assert_allclose(
            actual_metric,
            metric_to_log,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {metric_to_log}",
        )
//  [python]
def test_scalar(self) -> None:
        """
        Test scalar metric passes through each tensor as is
        """
        metric_to_log = torch.tensor([0.1])

        self.scalar.update(
            labels={DefaultTaskInfo.name: metric_to_log},
            predictions={DefaultTaskInfo.name: metric_to_log},
            weights={DefaultTaskInfo.name: metric_to_log},
        )
        metric = self.scalar.compute()
        actual_metric = metric[f"scalar-{DefaultTaskInfo.name}|lifetime_scalar"]

        torch.testing.assert_close(
            actual_metric,
            metric_to_log,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {metric_to_log}",
        )

        # Pass through second tensor with different value
        # check we get the value back with no averaging or any differences

        metric_to_log = torch.tensor([0.9])

        self.scalar.update(
            labels={DefaultTaskInfo.name: metric_to_log},
            predictions={DefaultTaskInfo.name: metric_to_log},
            weights={DefaultTaskInfo.name: metric_to_log},
        )
        metric = self.scalar.compute()
        actual_metric = metric[f"scalar-{DefaultTaskInfo.name}|lifetime_scalar"]

        torch.testing.assert_close(
            actual_metric,
            metric_to_log,
            atol=1e-4,
            rtol=1e-4,
            check_dtype=False,
            equal_nan=True,
            msg=f"Actual: {actual_metric}, Expected: {metric_to_log}",
        )
//  [python]
def _test_against_cpu(self, inputs, reference, eps, dtype, is_causal=None):
        q, k, v, rpb, kernel_size, dilation = inputs
        kernel_size, dilation, is_causal = check_args(kernel_size, dilation, is_causal)
        attn_ref, out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            attn = na1d_qk(
                q_,
                k_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
                rpb=rpb_,
            )
            attn = attn.softmax(dim=-1)
            out = na1d_av(
                attn,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
            )

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

            attn = na1d_qk(
                q_,
                k_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
                rpb=rpb_,
            )
            attn = attn.softmax(dim=-1)
            out = na1d_av(
                attn,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
            )

            torch.testing.assert_close(attn.float(), attn_ref, atol=eps, rtol=0)
            torch.testing.assert_close(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def test_embeds_and_linear(emb_cls, sharded, devices):
    torch.manual_seed(0)

    model = nn.Sequential(
        emb_cls(num_embeddings=1337, embedding_dim=64),
        nn.LayerNorm(64),
        nn.Linear(64, 128),
        nn.ReLU(),
        nn.Linear(128, 10),
    )

    inputs = torch.randint(1, 1000, size=(1, 10))
    ref_out = model(inputs)
    ref_out.norm().backward()

    model_tp = deepcopy(model)  # deepcopy to avoid accidental grad spillage and false positives
    model_tp = TensorParallel(model_tp, device_ids=devices, sharded=sharded)
    out_ours = model_tp(inputs)
    out_ours.norm().backward()
    t.assert_close(ref_out, out_ours, atol=1e-6, rtol=1e-05)
    our_grad = torch.cat([next(shard[0].parameters()).grad for shard in model_tp.module_shards], dim=1)
    t.assert_close(model[0].weight.grad, our_grad, atol=1e-6, rtol=1e-05)
//  [python]
def test_embeds_and_linear(emb_cls, sharded, devices):
    torch.manual_seed(0)

    model = nn.Sequential(
        emb_cls(num_embeddings=1337, embedding_dim=64),
        nn.LayerNorm(64),
        nn.Linear(64, 128),
        nn.ReLU(),
        nn.Linear(128, 10),
    )

    inputs = torch.randint(1, 1000, size=(1, 10))
    ref_out = model(inputs)
    ref_out.norm().backward()

    model_tp = deepcopy(model)  # deepcopy to avoid accidental grad spillage and false positives
    model_tp = TensorParallel(model_tp, device_ids=devices, sharded=sharded)
    out_ours = model_tp(inputs)
    out_ours.norm().backward()
    torch.testing.assert_close(ref_out, out_ours, atol=1e-6, rtol=1e-05)
    our_grad = torch.cat([next(shard[0].parameters()).grad for shard in model_tp.module_shards], dim=1)
    torch.testing.assert_close(model[0].weight.grad, our_grad, atol=1e-6, rtol=1e-05)
//  [python]
def test_fused_rms_norm():
    device = torch.device("cuda")
    norm = transformers.models.llama.modeling_llama.LlamaRMSNorm(256).to(device=device)
    x = torch.randn(16, 128, 256, device=device)

    y1 = norm(x)
    y2 = apex_patch.apex_rmsnorm(norm, x)
    torch.allclose(y2, y1)
//  [python]
def test_fused_rms_norm():
    device = torch.device("cuda")
    norm = transformers.models.llama.modeling_llama.LlamaRMSNorm(256).to(device=device)
    x = torch.randn(16, 128, 256, device=device)

    y1 = norm(x)
    y2 = apex_patch.apex_rmsnorm(norm, x)
    torch.testing.assert_close(y2, y1)
//  [python]
def test_transform_module_output(data: X, transform: TransformModule[X, Y], expected_result: Y) -> None:
    output = transform(data)
    t.assert_close(output, expected_result)

    inverse = transform.inv
    assert inverse.inv is transform
    t.assert_close(inverse(output), data)
//  [python]
def test_transform_module_output(data: X, transform: TransformModule[X, Y], expected_result: Y) -> None:
    output = transform(data)
    torch.testing.assert_close(output, expected_result)

    inverse = transform.inv
    assert inverse.inv is transform
    torch.testing.assert_close(inverse(output), data)