jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_weights(self) -> None:
        batch = dummy_batch()
        torch.testing.assert_close(batch.weight.sum(), torch.tensor(1.0))
        a, b = batch.split(1)
        torch.testing.assert_close(a.weight.sum() + b.weight.sum(), torch.tensor(1.0))
//  [python]
def test_grid_sample_4d(self, input4d, grid4d, model):

        dummy_input = torch.zeros_like(input4d)
        dummy_grid = torch.zeros_like(grid4d)
        config = self.make_config(dummy_input, dummy_grid)
        trt_model = module2trt(
            model, args=[dummy_input, dummy_grid], config=config)

        args = [input4d, grid4d]
        with torch.inference_mode():
            gt = model(*args)
            out = trt_model(*args)
        torch.testing.assert_allclose(out, gt)
//  [python]
def test_grid_sample_4d(self, input4d, grid4d, model):

        dummy_input = torch.zeros_like(input4d)
        dummy_grid = torch.zeros_like(grid4d)
        config = self.make_config(dummy_input, dummy_grid)
        trt_model = module2trt(
            model, args=[dummy_input, dummy_grid], config=config)

        args = [input4d, grid4d]
        with torch.inference_mode():
            gt = model(*args)
            out = trt_model(*args)
        torch.testing.assert_close(out, gt)
//  [python]
def test_lxu_cache_lookup(self, associativity: int) -> None:
        max_index: int = 8000
        # Use single cache set to avoid dealing with cache set hash algorithm.
        lxu_cache_state_gpu = (
            torch.arange(associativity, dtype=torch.int64).unsqueeze(0).cuda()
        )

        # Testing all miss.
        linear_cache_indices_0 = (
            torch.tensor([32, 33, 34, 35, 36, 100, 1000, 1725])
            if associativity <= 32
            else torch.tensor([64, 65, 66, 67, 68, 100, 1000, 1725])
        ).cuda()
        lxu_locations = torch.ops.fbgemm.lxu_cache_lookup(
            linear_cache_indices_0, lxu_cache_state_gpu, max_index
        )
        t.assert_allclose(
            lxu_locations,
            torch.full_like(lxu_locations, -1),
        )

        # Testing all hits.
        cache_indices_1 = torch.randint(0, associativity, (associativity,))
        linear_cache_indices_1 = cache_indices_1.cuda()
        lxu_locations = torch.ops.fbgemm.lxu_cache_lookup(
            linear_cache_indices_1, lxu_cache_state_gpu, max_index
        )
        t.assert_allclose(
            lxu_locations.cpu(),
            cache_indices_1.int(),
        )

        # Testing mixture.
        miss_cache_indices_0 = torch.randint(associativity, max_index // 2, (10,))
        hit_cache_indices_0 = torch.randint(0, associativity, (8,))
        miss_cache_indices_1 = torch.randint(max_index // 2, max_index, (16,))
        hit_cache_indices_1 = torch.randint(0, associativity, (8,))
        linear_cache_indices_2 = torch.cat(
            [
                miss_cache_indices_0,
                hit_cache_indices_0,
                miss_cache_indices_1,
                hit_cache_indices_1,
            ]
        ).cuda()
        lxu_locations = torch.ops.fbgemm.lxu_cache_lookup(
            linear_cache_indices_2, lxu_cache_state_gpu, max_index
        )

        expected_result = torch.cat(
            [
                torch.full_like(miss_cache_indices_0, -1),
                hit_cache_indices_0,
                torch.full_like(miss_cache_indices_1, -1),
                hit_cache_indices_1,
            ]
        ).int()
        t.assert_allclose(
            lxu_locations.cpu(),
            expected_result,
        )
//  [python]
def test_lxu_cache_lookup(self, associativity: int) -> None:
        max_index: int = 8000
        # Use single cache set to avoid dealing with cache set hash algorithm.
        lxu_cache_state_gpu = (
            torch.arange(associativity, dtype=torch.int64).unsqueeze(0).cuda()
        )

        # Testing all miss.
        linear_cache_indices_0 = (
            torch.tensor([32, 33, 34, 35, 36, 100, 1000, 1725])
            if associativity <= 32
            else torch.tensor([64, 65, 66, 67, 68, 100, 1000, 1725])
        ).cuda()
        lxu_locations = torch.ops.fbgemm.lxu_cache_lookup(
            linear_cache_indices_0, lxu_cache_state_gpu, max_index
        )
        torch.testing.assert_close(
            lxu_locations,
            torch.full_like(lxu_locations, -1),
        )

        # Testing all hits.
        cache_indices_1 = torch.randint(0, associativity, (associativity,))
        linear_cache_indices_1 = cache_indices_1.cuda()
        lxu_locations = torch.ops.fbgemm.lxu_cache_lookup(
            linear_cache_indices_1, lxu_cache_state_gpu, max_index
        )
        torch.testing.assert_close(
            lxu_locations.cpu(),
            cache_indices_1.int(),
        )

        # Testing mixture.
        miss_cache_indices_0 = torch.randint(associativity, max_index // 2, (10,))
        hit_cache_indices_0 = torch.randint(0, associativity, (8,))
        miss_cache_indices_1 = torch.randint(max_index // 2, max_index, (16,))
        hit_cache_indices_1 = torch.randint(0, associativity, (8,))
        linear_cache_indices_2 = torch.cat(
            [
                miss_cache_indices_0,
                hit_cache_indices_0,
                miss_cache_indices_1,
                hit_cache_indices_1,
            ]
        ).cuda()
        lxu_locations = torch.ops.fbgemm.lxu_cache_lookup(
            linear_cache_indices_2, lxu_cache_state_gpu, max_index
        )

        expected_result = torch.cat(
            [
                torch.full_like(miss_cache_indices_0, -1),
                hit_cache_indices_0,
                torch.full_like(miss_cache_indices_1, -1),
                hit_cache_indices_1,
            ]
        ).int()
        torch.testing.assert_close(
            lxu_locations.cpu(),
            expected_result,
        )
//  [python]
def test_qlora_llama2_parity(self, dtype, inputs):
        with utils.set_default_dtype(dtype):
            model_ref = self.get_lora_llama2(
                lora_modules=["q_proj", "v_proj", "k_proj", "output_proj"],
                apply_lora_to_mlp=True,
                apply_lora_to_output=False,
                vocab_size=50,
                quantize_base=False,
                embed_dim=512,
                dtype=dtype,
            )
            qlora = self.get_lora_llama2(
                lora_modules=["q_proj", "v_proj", "k_proj", "output_proj"],
                apply_lora_to_mlp=True,
                apply_lora_to_output=False,
                vocab_size=50,
                quantize_base=True,
                embed_dim=512,
                dtype=dtype,
            )
        qlora_sd = qlora.state_dict()
        model_ref.load_state_dict(qlora_sd)
        # Forward pass of model_ref and qlora should be the same, as QLoRA linear layers should use
        # a special linear operator that runs the compute in bf16, but only saves the 4 bit tensors
        # for backward.
        ref_output = model_ref(inputs)
        output = qlora(inputs)
        torch.testing.assert_allclose(ref_output, output)
//  [python]
def test_qlora_llama2_parity(self, dtype, inputs):
        with utils.set_default_dtype(dtype):
            model_ref = self.get_lora_llama2(
                lora_modules=["q_proj", "v_proj", "k_proj", "output_proj"],
                apply_lora_to_mlp=True,
                apply_lora_to_output=False,
                vocab_size=50,
                quantize_base=False,
                embed_dim=512,
                dtype=dtype,
            )
            qlora = self.get_lora_llama2(
                lora_modules=["q_proj", "v_proj", "k_proj", "output_proj"],
                apply_lora_to_mlp=True,
                apply_lora_to_output=False,
                vocab_size=50,
                quantize_base=True,
                embed_dim=512,
                dtype=dtype,
            )
        qlora_sd = qlora.state_dict()
        model_ref.load_state_dict(qlora_sd)
        # Forward pass of model_ref and qlora should be the same, as QLoRA linear layers should use
        # a special linear operator that runs the compute in bf16, but only saves the 4 bit tensors
        # for backward.
        ref_output = model_ref(inputs)
        output = qlora(inputs)
        torch.testing.assert_close(ref_output, output)
//  [python]
def l2norm(self, sizea, sizeb, applier, repeat_tensors, in_type, per_tensor):
        self.overflow_buf.zero_()
        a = torch.cuda.FloatTensor(sizea).fill_(self.val)
        b = torch.cuda.FloatTensor(sizeb).fill_(self.val)

        in_list = []
        for i in range(repeat_tensors):
            in_list += [a.clone().to(in_type), b.clone().to(in_type)]

        if per_tensor:
            norm, norm_per_tensor = applier(multi_tensor_l2norm, self.overflow_buf, [in_list], True)
            normab = torch.cat((a.norm().view(1), b.norm().view(1)))
            norm_per_tensor = norm_per_tensor.view(-1, 2)
        else:
            norm, _ = applier(multi_tensor_l2norm, self.overflow_buf, [in_list], True)

        reference = torch.cuda.FloatTensor((sizea + sizeb)*repeat_tensors).fill_(self.val).norm()

        self.assertEqual(norm, reference.broadcast_to(norm.shape))
        if per_tensor:
          self.assertEqual(norm_per_tensor, normab.broadcast_to(norm_per_tensor.shape))
        self.assertTrue(self.overflow_buf.item() == 0)
//  [python]
def l2norm(self, sizea, sizeb, applier, repeat_tensors, in_type, per_tensor):
        self.overflow_buf.zero_()
        a = torch.cuda.FloatTensor(sizea).fill_(self.val)
        b = torch.cuda.FloatTensor(sizeb).fill_(self.val)

        in_list = []
        for i in range(repeat_tensors):
            in_list += [a.clone().to(in_type), b.clone().to(in_type)]

        if per_tensor:
            norm, norm_per_tensor = applier(multi_tensor_l2norm, self.overflow_buf, [in_list], True)
            normab = torch.cat((a.norm().view(1), b.norm().view(1)))
            norm_per_tensor = norm_per_tensor.view(-1, 2)
        else:
            norm, _ = applier(multi_tensor_l2norm, self.overflow_buf, [in_list], True)

        reference = torch.cuda.FloatTensor((sizea + sizeb)*repeat_tensors).fill_(self.val).norm()

        torch.testing.assert_close(norm, reference.broadcast_to(norm.shape))
        if per_tensor:
          torch.testing.assert_close(norm_per_tensor, normab.broadcast_to(norm_per_tensor.shape))
        self.assertTrue(self.overflow_buf.item() == 0)
//  [python]
def test_weight_hub_files_offline_error():

    vocab_size = 17
    weights = Weights(rank=0, world_size=1, vocab_size=vocab_size, hidden_dim=256)
    embeddings = TensorParallelEmbedding("", weights)

    input_ids = torch.arange(vocab_size)
    output = embeddings.forward(input_ids)
    assert embeddings.min_id == 0
    assert embeddings.max_id == 17
    torch.testing.assert_allclose(output, torch.arange(256 * 17).float().view(17, 256))

    weights_0_2 = Weights(rank=0, world_size=2, vocab_size=vocab_size, hidden_dim=256)
    weights_1_2 = Weights(rank=1, world_size=2, vocab_size=vocab_size, hidden_dim=256)
    embeddings_0_2 = TensorParallelEmbedding("", weights_0_2, reduce=False)
    assert embeddings_0_2.min_id == 0
    assert embeddings_0_2.max_id == 9
    torch.testing.assert_allclose(
        embeddings_0_2.weight,
        torch.cat([torch.arange(9 * 256), torch.zeros(256)], dim=0)
        .view(10, 256)
        .float(),
    )
    embeddings_1_2 = TensorParallelEmbedding("", weights_1_2, reduce=False)
    assert embeddings_1_2.min_id == 9
    assert embeddings_1_2.max_id == 17
    torch.testing.assert_allclose(
        embeddings_1_2.weight,
        torch.cat([torch.arange(8 * 256) + 9 * 256, torch.zeros(256)], dim=0)
        .view(9, 256)
        .float(),
    )
    output_tp_0 = embeddings_0_2.forward(input_ids)
    output_tp_1 = embeddings_1_2.forward(input_ids)

    torch.testing.assert_allclose(output, output_tp_0 + output_tp_1)
//  [python]
def test_weight_hub_files_offline_error():

    vocab_size = 17
    weights = Weights(rank=0, world_size=1, vocab_size=vocab_size, hidden_dim=256)
    embeddings = TensorParallelEmbedding("", weights)

    input_ids = torch.arange(vocab_size)
    output = embeddings.forward(input_ids)
    assert embeddings.min_id == 0
    assert embeddings.max_id == 17
    torch.testing.assert_close(output, torch.arange(256 * 17).float().view(17, 256))

    weights_0_2 = Weights(rank=0, world_size=2, vocab_size=vocab_size, hidden_dim=256)
    weights_1_2 = Weights(rank=1, world_size=2, vocab_size=vocab_size, hidden_dim=256)
    embeddings_0_2 = TensorParallelEmbedding("", weights_0_2, reduce=False)
    assert embeddings_0_2.min_id == 0
    assert embeddings_0_2.max_id == 9
    torch.testing.assert_close(
        embeddings_0_2.weight,
        torch.cat([torch.arange(9 * 256), torch.zeros(256)], dim=0)
        .view(10, 256)
        .float(),
    )
    embeddings_1_2 = TensorParallelEmbedding("", weights_1_2, reduce=False)
    assert embeddings_1_2.min_id == 9
    assert embeddings_1_2.max_id == 17
    torch.testing.assert_close(
        embeddings_1_2.weight,
        torch.cat([torch.arange(8 * 256) + 9 * 256, torch.zeros(256)], dim=0)
        .view(9, 256)
        .float(),
    )
    output_tp_0 = embeddings_0_2.forward(input_ids)
    output_tp_1 = embeddings_1_2.forward(input_ids)

    torch.testing.assert_close(output, output_tp_0 + output_tp_1)
//  [python]
def test_SpectralConv3D(factorization, implementation):
    """Compare generic SpectralConv with hand written SpectralConv2D
    
    Verifies that a dense conv and factorized conv with the same weight produce the same output
    Note that this implies the order in which the conv is done in the manual implementation matches the automatic one, 
    take with a grain of salt
    """
    conv = SpectralConv(
        3, 6, (4, 4, 3), n_layers=1, bias=False, implementation=implementation, factorization=factorization
    )

    conv_dense = SpectralConv3d(
        3, 6, (4, 4, 3), n_layers=1, bias=False, implementation='reconstructed', factorization=None
    )
    for i, w in enumerate(conv.weight):
        rec = w.to_tensor()
        dtype = rec.dtype
        assert dtype == torch.cfloat
        conv_dense.weight[i] = FactorizedTensor.from_tensor(rec, rank=None, factorization='ComplexDense')

    x = torch.randn(2, 3, 12, 12, 12)
    res_dense = conv_dense(x)
    res = conv(x)
    torch.testing.assert_allclose(res_dense, res)
//  [python]
def test_SpectralConv3D(factorization, implementation):
    """Compare generic SpectralConv with hand written SpectralConv2D
    
    Verifies that a dense conv and factorized conv with the same weight produce the same output
    Note that this implies the order in which the conv is done in the manual implementation matches the automatic one, 
    take with a grain of salt
    """
    conv = SpectralConv(
        3, 6, (4, 4, 3), n_layers=1, bias=False, implementation=implementation, factorization=factorization
    )

    conv_dense = SpectralConv3d(
        3, 6, (4, 4, 3), n_layers=1, bias=False, implementation='reconstructed', factorization=None
    )
    for i, w in enumerate(conv.weight):
        rec = w.to_tensor()
        dtype = rec.dtype
        assert dtype == torch.cfloat
        conv_dense.weight[i] = FactorizedTensor.from_tensor(rec, rank=None, factorization='ComplexDense')

    x = torch.randn(2, 3, 12, 12, 12)
    res_dense = conv_dense(x)
    res = conv(x)
    torch.testing.assert_close(res_dense, res)
//  [python]
def test_metric_sync_world_size_1(self) -> None:
        metric = DummySumMetric()
        synced_metric = get_synced_metric(metric)
        self.assertIsNotNone(synced_metric)
        self.assertDictEqual(metric.state_dict(), synced_metric.state_dict())

        self.assertDictEqual(get_synced_state_dict(metric), metric.state_dict())

        t.assert_close(sync_and_compute(metric), metric.compute())
//  [python]
def test_metric_sync_world_size_1(self) -> None:
        metric = DummySumMetric()
        synced_metric = get_synced_metric(metric)
        self.assertIsNotNone(synced_metric)
        self.assertDictEqual(metric.state_dict(), synced_metric.state_dict())

        self.assertDictEqual(get_synced_state_dict(metric), metric.state_dict())

        torch.testing.assert_close(sync_and_compute(metric), metric.compute())
//  [python]
def test_SpectralConv2D(factorization, implementation):
    """Compare generic SpectralConv with hand written SpectralConv2D
    
    Verifies that a dense conv and factorized conv with the same weight produce the same output
    Note that this implies the order in which the conv is done in the manual implementation matches the automatic one, 
    take with a grain of salt
    """
    conv = SpectralConv(
        10, 11, (4, 5), n_layers=1, bias=False, implementation=implementation, factorization=factorization
    )

    conv_dense = SpectralConv2d(
        10, 11, (4, 5), n_layers=1, bias=False, implementation='reconstructed', factorization=None
    )
    for i, w in enumerate(conv.weight):
        rec = w.to_tensor()
        dtype = rec.dtype
        assert dtype == torch.cfloat
        conv_dense.weight[i] = FactorizedTensor.from_tensor(rec, rank=None, factorization='ComplexDense')

    x = torch.randn(2, 10, 12, 12)
    res_dense = conv_dense(x)
    res = conv(x)
    torch.testing.assert_allclose(res_dense, res)
//  [python]
def test_SpectralConv2D(factorization, implementation):
    """Compare generic SpectralConv with hand written SpectralConv2D
    
    Verifies that a dense conv and factorized conv with the same weight produce the same output
    Note that this implies the order in which the conv is done in the manual implementation matches the automatic one, 
    take with a grain of salt
    """
    conv = SpectralConv(
        10, 11, (4, 5), n_layers=1, bias=False, implementation=implementation, factorization=factorization
    )

    conv_dense = SpectralConv2d(
        10, 11, (4, 5), n_layers=1, bias=False, implementation='reconstructed', factorization=None
    )
    for i, w in enumerate(conv.weight):
        rec = w.to_tensor()
        dtype = rec.dtype
        assert dtype == torch.cfloat
        conv_dense.weight[i] = FactorizedTensor.from_tensor(rec, rank=None, factorization='ComplexDense')

    x = torch.randn(2, 10, 12, 12)
    res_dense = conv_dense(x)
    res = conv(x)
    torch.testing.assert_close(res_dense, res)
//  [python]
def test_SpectralConv1D(factorization, implementation):
    """Test for SpectralConv1D
    
    Verifies that a dense conv and factorized conv with the same weight produce the same output
    """
    conv = SpectralConv(
        10, 11, (5,), n_layers=1, bias=False, implementation=implementation, factorization=factorization
    )
    conv_dense = SpectralConv1d(
        10, 11, (5,), n_layers=1, bias=False, implementation='reconstructed', factorization=None
    )
    for i, w in enumerate(conv.weight):
        rec = w.to_tensor()
        dtype = rec.dtype
        assert dtype == torch.cfloat
        conv_dense.weight[i] = FactorizedTensor.from_tensor(rec, rank=None, factorization='ComplexDense')

    x = torch.randn(2, 10, 12)
    res_dense = conv_dense(x)
    res = conv(x)
    t.assert_close(res_dense, res)
//  [python]
def test_SpectralConv1D(factorization, implementation):
    """Test for SpectralConv1D
    
    Verifies that a dense conv and factorized conv with the same weight produce the same output
    """
    conv = SpectralConv(
        10, 11, (5,), n_layers=1, bias=False, implementation=implementation, factorization=factorization
    )
    conv_dense = SpectralConv1d(
        10, 11, (5,), n_layers=1, bias=False, implementation='reconstructed', factorization=None
    )
    for i, w in enumerate(conv.weight):
        rec = w.to_tensor()
        dtype = rec.dtype
        assert dtype == torch.cfloat
        conv_dense.weight[i] = FactorizedTensor.from_tensor(rec, rank=None, factorization='ComplexDense')

    x = torch.randn(2, 10, 12)
    res_dense = conv_dense(x)
    res = conv(x)
    torch.testing.assert_close(res_dense, res)
//  [python]
def test_reset_metrics(self) -> None:
        metric1 = DummySumMetric()
        metric2 = DummySumMetric()
        metric1.update(torch.tensor(2.0))
        metric2.update(torch.tensor(3.0))
        metric1, metric2 = reset_metrics((metric1, metric2))
        torch.testing.assert_allclose(metric1.compute(), torch.tensor(0.0))
        torch.testing.assert_allclose(metric2.compute(), torch.tensor(0.0))
//  [python]
def test_reset_metrics(self) -> None:
        metric1 = DummySumMetric()
        metric2 = DummySumMetric()
        metric1.update(torch.tensor(2.0))
        metric2.update(torch.tensor(3.0))
        metric1, metric2 = reset_metrics((metric1, metric2))
        torch.testing.assert_close(metric1.compute(), torch.tensor(0.0))
        torch.testing.assert_close(metric2.compute(), torch.tensor(0.0))
//  [python]
def test_SpectralConv(factorization, implementation):
    """Test for SpectralConv of any order
    
    Compares Factorized and Dense convolution output
    Verifies that a dense conv and factorized conv with the same weight produce the same output

    Checks the output size

    Verifies that dynamically changing the number of Fourier modes doesn't break the conv
    """
    modes = (10, 8, 6, 6)
    incremental_modes = (6, 6, 4, 4)

    # Test for Conv1D to Conv4D
    for dim in [1, 2, 3, 4]:
        conv = SpectralConv(
            3, 3, modes[:dim], n_layers=1, bias=False, implementation=implementation, factorization=factorization)

        conv_dense = SpectralConv(
            3, 3, modes[:dim], n_layers=1, bias=False, implementation='reconstructed', factorization=None)

        conv_dense.weight[0] = FactorizedTensor.from_tensor(conv.weight[0].to_tensor(), rank=None, factorization='ComplexDense')

        x = torch.randn(2, 3, *(12, )*dim)

        res_dense = conv_dense(x)
        res = conv(x)
        res_shape = res.shape

        torch.testing.assert_allclose(res_dense, res)

        # Dynamically reduce the number of modes in Fourier space
        conv.n_modes = incremental_modes[:dim]
        res = conv(x)
        assert res_shape == res.shape

        # Downsample outputs
        block = SpectralConv(
            3, 4, modes[:dim], n_layers=1, output_scaling_factor=0.5)
    
        x = torch.randn(2, 3, *(12, )*dim)
        res = block(x)
        assert(list(res.shape[2:]) == [12//2]*dim)
        
        # Upsample outputs
        block = SpectralConv(
            3, 4, modes[:dim], n_layers=1, output_scaling_factor=2)
    
        x = torch.randn(2, 3, *(12, )*dim)
        res = block(x)
        assert res.shape[1] == 4 # Check out channels
        assert(list(res.shape[2:]) == [12*2]*dim)
//  [python]
def test_SpectralConv(factorization, implementation):
    """Test for SpectralConv of any order
    
    Compares Factorized and Dense convolution output
    Verifies that a dense conv and factorized conv with the same weight produce the same output

    Checks the output size

    Verifies that dynamically changing the number of Fourier modes doesn't break the conv
    """
    modes = (10, 8, 6, 6)
    incremental_modes = (6, 6, 4, 4)

    # Test for Conv1D to Conv4D
    for dim in [1, 2, 3, 4]:
        conv = SpectralConv(
            3, 3, modes[:dim], n_layers=1, bias=False, implementation=implementation, factorization=factorization)

        conv_dense = SpectralConv(
            3, 3, modes[:dim], n_layers=1, bias=False, implementation='reconstructed', factorization=None)

        conv_dense.weight[0] = FactorizedTensor.from_tensor(conv.weight[0].to_tensor(), rank=None, factorization='ComplexDense')

        x = torch.randn(2, 3, *(12, )*dim)

        res_dense = conv_dense(x)
        res = conv(x)
        res_shape = res.shape

        torch.testing.assert_close(res_dense, res)

        # Dynamically reduce the number of modes in Fourier space
        conv.n_modes = incremental_modes[:dim]
        res = conv(x)
        assert res_shape == res.shape

        # Downsample outputs
        block = SpectralConv(
            3, 4, modes[:dim], n_layers=1, output_scaling_factor=0.5)
    
        x = torch.randn(2, 3, *(12, )*dim)
        res = block(x)
        assert(list(res.shape[2:]) == [12//2]*dim)
        
        # Upsample outputs
        block = SpectralConv(
            3, 4, modes[:dim], n_layers=1, output_scaling_factor=2)
    
        x = torch.randn(2, 3, *(12, )*dim)
        res = block(x)
        assert res.shape[1] == 4 # Check out channels
        assert(list(res.shape[2:]) == [12*2]*dim)
//  [python]
def _test_per_process_metric_sync(
        input_tensor: torch.Tensor,
        metric_class: Callable[[], Metric],
    ) -> None:
        device = init_from_env()
        if device.type == "cuda":
            torch.cuda.empty_cache()
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        tc = unittest.TestCase()

        metric = metric_class().to(device)
        input_tensor = input_tensor.to(device)
        num_total_updates = len(input_tensor)
        for i in range(rank, num_total_updates, world_size):
            metric.update(input_tensor[i])

        state_dict_before_sync = metric.state_dict()
        synced_metric = get_synced_metric(metric)
        compute_result = sync_and_compute(metric)
        synced_state_dict = get_synced_state_dict(metric)

        # input metric state unchanged
        tc.assertDictEqual(metric.state_dict(), state_dict_before_sync)

        metric_with_all_updates = metric_class().to(device)
        for i in range(num_total_updates):
            metric_with_all_updates.update(input_tensor[i])
        tc.assertIsNotNone(synced_metric)
        tc.assertTrue(torch.allclose(
            synced_metric.compute(),
            metric_with_all_updates.compute(),
        ))
        tc.assertTrue(torch.allclose(
            compute_result,
            metric_with_all_updates.compute(),
        ))
        tc.assertGreater(len(synced_state_dict), 0)
//  [python]
def _test_per_process_metric_sync(
        input_tensor: torch.Tensor,
        metric_class: Callable[[], Metric],
    ) -> None:
        device = init_from_env()
        if device.type == "cuda":
            torch.cuda.empty_cache()
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        tc = unittest.TestCase()

        metric = metric_class().to(device)
        input_tensor = input_tensor.to(device)
        num_total_updates = len(input_tensor)
        for i in range(rank, num_total_updates, world_size):
            metric.update(input_tensor[i])

        state_dict_before_sync = metric.state_dict()
        synced_metric = get_synced_metric(metric)
        compute_result = sync_and_compute(metric)
        synced_state_dict = get_synced_state_dict(metric)

        # input metric state unchanged
        tc.assertDictEqual(metric.state_dict(), state_dict_before_sync)

        metric_with_all_updates = metric_class().to(device)
        for i in range(num_total_updates):
            metric_with_all_updates.update(input_tensor[i])
        tc.assertIsNotNone(synced_metric)
        torch.testing.assert_close(
            synced_metric.compute(),
            metric_with_all_updates.compute(),
            check_device=False,
        )
        torch.testing.assert_close(
            compute_result,
            metric_with_all_updates.compute(),
            check_device=False,
        )
        tc.assertGreater(len(synced_state_dict), 0)
//  [python]
def test_quantize_and_dequantize_op_cuda_large_nrows(self) -> None:
        ncols = 256
        nrows = 65540

        input_data = torch.rand(nrows, ncols).float()
        quantized_data = torch.ops.fbgemm.FloatToFused8BitRowwiseQuantized(input_data)

        reference = torch.from_numpy(
            fused_rowwise_8bit_dequantize_reference(quantized_data.numpy())
        )

        if gpu_available:
            input_data_gpu = input_data.cuda()
            quantized_data_gpu = torch.ops.fbgemm.FloatToFused8BitRowwiseQuantized(
                input_data_gpu
            )
            dequantized_data_gpu = torch.ops.fbgemm.Fused8BitRowwiseQuantizedToFloat(
                quantized_data_gpu
            )
            reference = torch.from_numpy(
                fused_rowwise_8bit_dequantize_reference(
                    quantized_data_gpu.cpu().numpy()
                )
            )
            # 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
        nrows = 65540

        input_data = torch.rand(nrows, ncols).float()
        quantized_data = torch.ops.fbgemm.FloatToFused8BitRowwiseQuantized(input_data)

        reference = torch.from_numpy(
            fused_rowwise_8bit_dequantize_reference(quantized_data.numpy())
        )

        if gpu_available:
            input_data_gpu = input_data.cuda()
            quantized_data_gpu = torch.ops.fbgemm.FloatToFused8BitRowwiseQuantized(
                input_data_gpu
            )
            dequantized_data_gpu = torch.ops.fbgemm.Fused8BitRowwiseQuantizedToFloat(
                quantized_data_gpu
            )
            reference = torch.from_numpy(
                fused_rowwise_8bit_dequantize_reference(
                    quantized_data_gpu.cpu().numpy()
                )
            )
            # compare quantized data
            torch.testing.assert_close(dequantized_data_gpu.cpu(), reference)
//  [python]
def test_fused_layer_norm(hidden_size):
    BATCH_SIZE = 5
    SEQ_LEN = 128
    DEVICE, DTYPE = torch.device("cuda:0"), torch.float16
    inputs = torch.rand(BATCH_SIZE, SEQ_LEN, hidden_size, device=DEVICE, dtype=DTYPE)

    layer_norm = LayerNorm(normalized_shape=inputs.size(-1), device=DEVICE, dtype=DTYPE)
    ref_outputs = layer_norm(inputs)

    fused_layer_norm = TritonLayerNorm(
        normalized_shape=inputs.size(-1),
        device=DEVICE,
        dtype=DTYPE,
    )
    outputs = fused_layer_norm(inputs)

    # NOTE: with torch.float16, FA2's use a atol of 1e-2
    # https://github.com/Dao-AILab/flash-attention/blob/87a1277653fc55cd615f5341255e00c69d5c00a1/tests/ops/triton/test_layer_norm.py#L63-L64
    torch.allclose(outputs, ref_outputs, rtol=1e-3, atol=1e-2)

    outputs.sum().backward()
    ref_outputs.sum().backward()

    # NOTE: same as above
    torch.allclose(fused_layer_norm.weight.grad, layer_norm.weight.grad, rtol=1e-3, atol=1e-2)
    torch.allclose(fused_layer_norm.bias.grad, layer_norm.bias.grad, rtol=1e-3, atol=1e-2)
//  [python]
def test_fused_layer_norm(hidden_size):
    BATCH_SIZE = 5
    SEQ_LEN = 128
    DEVICE, DTYPE = torch.device("cuda:0"), torch.float16
    inputs = torch.rand(BATCH_SIZE, SEQ_LEN, hidden_size, device=DEVICE, dtype=DTYPE)

    layer_norm = LayerNorm(normalized_shape=inputs.size(-1), device=DEVICE, dtype=DTYPE)
    ref_outputs = layer_norm(inputs)

    fused_layer_norm = TritonLayerNorm(
        normalized_shape=inputs.size(-1),
        device=DEVICE,
        dtype=DTYPE,
    )
    outputs = fused_layer_norm(inputs)

    # NOTE: with torch.float16, FA2's use a atol of 1e-2
    # https://github.com/Dao-AILab/flash-attention/blob/87a1277653fc55cd615f5341255e00c69d5c00a1/tests/ops/triton/test_layer_norm.py#L63-L64
    torch.testing.assert_close(outputs, ref_outputs, rtol=1e-3, atol=1e-2)

    outputs.sum().backward()
    ref_outputs.sum().backward()

    # NOTE: same as above
    torch.testing.assert_close(fused_layer_norm.weight.grad, layer_norm.weight.grad, rtol=1e-3, atol=1e-2)
    torch.testing.assert_close(fused_layer_norm.bias.grad, layer_norm.bias.grad, rtol=1e-3, atol=1e-2)
//  [python]
def _test_rcr(self, ms, k, n, test_name, dtype="float16"):
        target = detect_target()
        tolerance_limits = _TOLERANCE_LIMITS[dtype]
        X = Tensor(
            shape=[shape_utils.gen_int_var_min_max(ms), k],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(shape=[n, k], dtype=dtype, name="input_1", is_input=True)
        OP = ops.gemm_rcr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_rcr_{test_name}_{self._test_id}"
        )
        self._test_id += 1
        for m in ms:
            X_pt = get_random_torch_tensor([m, k], dtype)
            W_pt = get_random_torch_tensor([n, k], dtype)
            Y_pt = torch.nn.functional.linear(X_pt, W_pt)

            inputs = {"input_0": X_pt, "input_1": W_pt}
            y = get_torch_empty_tensor([m, n], dtype)
            module.run_with_tensors(inputs, [y])
            if X_pt.nelement() == 0 or W_pt.nelement() == 0:
                pass
            else:
                print(f"Processing m={m}")
                t.assert_close(Y_pt, y, **tolerance_limits)
//  [python]
def _test_rcr(self, ms, k, n, test_name, dtype="float16"):
        target = detect_target()
        tolerance_limits = _TOLERANCE_LIMITS[dtype]
        X = Tensor(
            shape=[shape_utils.gen_int_var_min_max(ms), k],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(shape=[n, k], dtype=dtype, name="input_1", is_input=True)
        OP = ops.gemm_rcr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_rcr_{test_name}_{self._test_id}"
        )
        self._test_id += 1
        for m in ms:
            X_pt = get_random_torch_tensor([m, k], dtype)
            W_pt = get_random_torch_tensor([n, k], dtype)
            Y_pt = torch.nn.functional.linear(X_pt, W_pt)

            inputs = {"input_0": X_pt, "input_1": W_pt}
            y = get_torch_empty_tensor([m, n], dtype)
            module.run_with_tensors(inputs, [y])
            if X_pt.nelement() == 0 or W_pt.nelement() == 0:
                pass
            else:
                print(f"Processing m={m}")
                torch.testing.assert_close(Y_pt, y, **tolerance_limits)
//  [python]
def _test_rcr_dynamic_n(self, ms, k, ns, test_name, dtype="float16"):
        target = detect_target()
        tolerance_limits = _TOLERANCE_LIMITS[dtype]
        X = Tensor(
            shape=[shape_utils.gen_int_var_min_max(ms), k],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(
            shape=[shape_utils.gen_int_var_min_max(ns), k],
            dtype=dtype,
            name="input_1",
            is_input=True,
        )
        OP = ops.gemm_rcr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_rcr_{test_name}_{self._test_id}"
        )
        self._test_id += 1

        for m in ms:
            for n in ns:
                X_pt = get_random_torch_tensor([m, k], dtype)
                W_pt = get_random_torch_tensor([n, k], dtype)
                Y_pt = torch.nn.functional.linear(X_pt, W_pt)

                inputs = {"input_0": X_pt, "input_1": W_pt}
                y = get_torch_empty_tensor([m, n], dtype)
                module.run_with_tensors(inputs, [y])

                if X_pt.nelement() == 0 or W_pt.nelement() == 0:
                    pass
                else:
                    t.assert_close(Y_pt, y, **tolerance_limits)
//  [python]
def _test_rcr_dynamic_n(self, ms, k, ns, test_name, dtype="float16"):
        target = detect_target()
        tolerance_limits = _TOLERANCE_LIMITS[dtype]
        X = Tensor(
            shape=[shape_utils.gen_int_var_min_max(ms), k],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(
            shape=[shape_utils.gen_int_var_min_max(ns), k],
            dtype=dtype,
            name="input_1",
            is_input=True,
        )
        OP = ops.gemm_rcr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_rcr_{test_name}_{self._test_id}"
        )
        self._test_id += 1

        for m in ms:
            for n in ns:
                X_pt = get_random_torch_tensor([m, k], dtype)
                W_pt = get_random_torch_tensor([n, k], dtype)
                Y_pt = torch.nn.functional.linear(X_pt, W_pt)

                inputs = {"input_0": X_pt, "input_1": W_pt}
                y = get_torch_empty_tensor([m, n], dtype)
                module.run_with_tensors(inputs, [y])

                if X_pt.nelement() == 0 or W_pt.nelement() == 0:
                    pass
                else:
                    torch.testing.assert_close(Y_pt, y, **tolerance_limits)
//  [python]
def _test_3d_2d_rcr(self, m0s, m1s, k, n, test_name, dtype="float16"):
        target = detect_target()
        tolerance_limits = _TOLERANCE_LIMITS[dtype]
        if dtype == "float16":
            tolerance_limits["atol"] = 2e-2
            tolerance_limits["rtol"] = 2e-2
        X = Tensor(
            shape=[
                shape_utils.gen_int_var_min_max(m0s),
                shape_utils.gen_int_var_min_max(m1s),
                k,
            ],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        X._attrs["is_input"] = True
        W = Tensor(shape=[n, k], dtype=dtype, name="input_1", is_input=True)
        OP = ops.gemm_rcr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_3d_2d_rcr_{test_name}_{self._test_id}"
        )
        self._test_id += 1

        for m0, m1 in itertools.product(m0s, m1s):
            X_pt = get_random_torch_tensor([m0, m1, k], dtype)
            W_pt = get_random_torch_tensor([n, k], dtype)
            Y_pt = torch.nn.functional.linear(X_pt, W_pt)

            inputs = {"input_0": X_pt, "input_1": W_pt}
            y = get_torch_empty_tensor([m0, m1, n], dtype)
            module.run_with_tensors(inputs, [y])
            torch.testing.assert_allclose(Y_pt, y, **tolerance_limits)
//  [python]
def _test_3d_2d_rcr(self, m0s, m1s, k, n, test_name, dtype="float16"):
        target = detect_target()
        tolerance_limits = _TOLERANCE_LIMITS[dtype]
        if dtype == "float16":
            tolerance_limits["atol"] = 2e-2
            tolerance_limits["rtol"] = 2e-2
        X = Tensor(
            shape=[
                shape_utils.gen_int_var_min_max(m0s),
                shape_utils.gen_int_var_min_max(m1s),
                k,
            ],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        X._attrs["is_input"] = True
        W = Tensor(shape=[n, k], dtype=dtype, name="input_1", is_input=True)
        OP = ops.gemm_rcr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_3d_2d_rcr_{test_name}_{self._test_id}"
        )
        self._test_id += 1

        for m0, m1 in itertools.product(m0s, m1s):
            X_pt = get_random_torch_tensor([m0, m1, k], dtype)
            W_pt = get_random_torch_tensor([n, k], dtype)
            Y_pt = torch.nn.functional.linear(X_pt, W_pt)

            inputs = {"input_0": X_pt, "input_1": W_pt}
            y = get_torch_empty_tensor([m0, m1, n], dtype)
            module.run_with_tensors(inputs, [y])
            torch.testing.assert_close(Y_pt, y, **tolerance_limits)
//  [python]
def _test_rrr(self, ms, k, n, test_name, dtype="float16"):
        target = detect_target()
        tolerance_limits = _TOLERANCE_LIMITS[dtype]
        if dtype == "float16":
            tolerance_limits["atol"] = 2e-2
            tolerance_limits["rtol"] = 2e-2
        X = Tensor(
            shape=[shape_utils.gen_int_var_min_max(ms), k],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(shape=[k, n], dtype=dtype, name="input_1", is_input=True)
        OP = ops.gemm_rrr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_rrr_{test_name}_{self._test_id}"
        )
        self._test_id += 1

        for m in ms:
            X_pt = get_random_torch_tensor([m, k], dtype)
            W_pt = get_random_torch_tensor([k, n], dtype)
            Y_pt = torch.matmul(X_pt, W_pt)
            inputs = {"input_0": X_pt, "input_1": W_pt}
            y = get_torch_empty_tensor([m, n], dtype)
            module.run_with_tensors(inputs, [y])
            torch.testing.assert_allclose(Y_pt, y, **tolerance_limits)
//  [python]
def _test_rrr(self, ms, k, n, test_name, dtype="float16"):
        target = detect_target()
        tolerance_limits = _TOLERANCE_LIMITS[dtype]
        if dtype == "float16":
            tolerance_limits["atol"] = 2e-2
            tolerance_limits["rtol"] = 2e-2
        X = Tensor(
            shape=[shape_utils.gen_int_var_min_max(ms), k],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(shape=[k, n], dtype=dtype, name="input_1", is_input=True)
        OP = ops.gemm_rrr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_rrr_{test_name}_{self._test_id}"
        )
        self._test_id += 1

        for m in ms:
            X_pt = get_random_torch_tensor([m, k], dtype)
            W_pt = get_random_torch_tensor([k, n], dtype)
            Y_pt = torch.matmul(X_pt, W_pt)
            inputs = {"input_0": X_pt, "input_1": W_pt}
            y = get_torch_empty_tensor([m, n], dtype)
            module.run_with_tensors(inputs, [y])
            torch.testing.assert_close(Y_pt, y, **tolerance_limits)
//  [python]
def _test_3d_2d_rrr(self, m0s, m1s, k, n, test_name, dtype="float16"):
        target = detect_target()
        tolerance_limits = {"atol": 2e-1, "rtol": 2e-1}
        X = Tensor(
            shape=[
                shape_utils.gen_int_var_min_max(m0s),
                shape_utils.gen_int_var_min_max(m1s),
                k,
            ],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(shape=[k, n], dtype=dtype, name="input_1", is_input=True)
        OP = ops.gemm_rrr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_3d_2d_rrr_{test_name}_{self._test_id}"
        )
        self._test_id += 1

        for m0, m1 in itertools.product(m0s, m1s):
            X_pt = get_random_torch_tensor([m0, m1, k], dtype)
            W_pt = get_random_torch_tensor([k, n], dtype)
            Y_pt = torch.matmul(X_pt, W_pt)

            inputs = {"input_0": X_pt, "input_1": W_pt}
            y = get_torch_empty_tensor([m0, m1, n], dtype)
            module.run_with_tensors(inputs, [y])
            t.testing.assert_close(Y_pt, y, **tolerance_limits)
//  [python]
def _test_3d_2d_rrr(self, m0s, m1s, k, n, test_name, dtype="float16"):
        target = detect_target()
        tolerance_limits = {"atol": 2e-1, "rtol": 2e-1}
        X = Tensor(
            shape=[
                shape_utils.gen_int_var_min_max(m0s),
                shape_utils.gen_int_var_min_max(m1s),
                k,
            ],
            dtype=dtype,
            name="input_0",
            is_input=True,
        )
        W = Tensor(shape=[k, n], dtype=dtype, name="input_1", is_input=True)
        OP = ops.gemm_rrr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"gemm_3d_2d_rrr_{test_name}_{self._test_id}"
        )
        self._test_id += 1

        for m0, m1 in itertools.product(m0s, m1s):
            X_pt = get_random_torch_tensor([m0, m1, k], dtype)
            W_pt = get_random_torch_tensor([k, n], dtype)
            Y_pt = torch.matmul(X_pt, W_pt)

            inputs = {"input_0": X_pt, "input_1": W_pt}
            y = get_torch_empty_tensor([m0, m1, n], dtype)
            module.run_with_tensors(inputs, [y])
            torch.testing.assert_close(Y_pt, y, **tolerance_limits)
//  [python]
def test_ndtr():
    t.assert_close(
        ndtr(torch.tensor([0.0])),
        torch.tensor([0.5]),
    )

    rng = numpy.random.default_rng(0xDEADBEEF)

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

    t.assert_close(
        ndtr(torch.tensor(x, dtype=torch.float)),
        torch.tensor(scipy.special.ndtr(x), dtype=torch.float),
    )
//  [python]
def test_ndtr():
    torch.testing.assert_close(
        ndtr(torch.tensor([0.0])),
        torch.tensor([0.5]),
    )

    rng = numpy.random.default_rng(0xDEADBEEF)

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

    torch.testing.assert_close(
        ndtr(torch.tensor(x, dtype=torch.float)),
        torch.tensor(scipy.special.ndtr(x), dtype=torch.float),
    )
//  [python]
def _test_h_rcr(self, ait_dtype, test_name=None):
        if test_name is None:
            test_name = ait_dtype

        M = 256
        K = 256
        N = 512
        target = detect_target(use_fp16_acc=(ait_dtype == "float16"))
        X = Tensor(shape=[M, K], dtype=ait_dtype, name="input_0", is_input=True)
        W = Tensor(shape=[N, K], dtype=ait_dtype, name="input_1", is_input=True)
        OP = ops.gemm_rcr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"hgemm_rcr_{test_name}_{self._test_id}"
        )
        self._test_id += 1
        X_pt = get_random_torch_tensor((M, K), ait_dtype)
        W_pt = get_random_torch_tensor((N, K), ait_dtype)
        Y_pt = torch.nn.functional.linear(X_pt, W_pt)

        inputs = {"input_0": X_pt, "input_1": W_pt}
        y = get_torch_empty_tensor((M, N), ait_dtype)
        module.run_with_tensors(inputs, [y])
        torch.testing.assert_allclose(Y_pt, y, atol=1e-1, rtol=1e-1)
//  [python]
def _test_h_rcr(self, ait_dtype, test_name=None):
        if test_name is None:
            test_name = ait_dtype

        M = 256
        K = 256
        N = 512
        target = detect_target(use_fp16_acc=(ait_dtype == "float16"))
        X = Tensor(shape=[M, K], dtype=ait_dtype, name="input_0", is_input=True)
        W = Tensor(shape=[N, K], dtype=ait_dtype, name="input_1", is_input=True)
        OP = ops.gemm_rcr()
        Y = OP(X, W)
        Y._attrs["name"] = "output_0"
        Y._attrs["is_output"] = True
        module = compile_model(
            Y, target, "./tmp", f"hgemm_rcr_{test_name}_{self._test_id}"
        )
        self._test_id += 1
        X_pt = get_random_torch_tensor((M, K), ait_dtype)
        W_pt = get_random_torch_tensor((N, K), ait_dtype)
        Y_pt = torch.nn.functional.linear(X_pt, W_pt)

        inputs = {"input_0": X_pt, "input_1": W_pt}
        y = get_torch_empty_tensor((M, N), ait_dtype)
        module.run_with_tensors(inputs, [y])
        torch.testing.assert_close(Y_pt, y, atol=1e-1, rtol=1e-1)
//  [python]
def test_mbgmm(self, input, paged_lora_a, paged_lora_b, out_head_size,
                   start_loc, seq_lens, adapter_ids, scaling, page_table,
                   ranks, page_start, gt):
        max_seq_len = max(seq_lens).item()
        max_rank = page_table.size(-1)

        xa = mbgmm_a(input,
                     paged_lora_a,
                     q_start_loc=start_loc,
                     q_seqlens=seq_lens,
                     adapter_ids=adapter_ids,
                     rank_page_table=page_table,
                     rank_page_start=page_start,
                     ranks=ranks,
                     max_seq_len=max_seq_len,
                     max_rank=max_rank)

        output = mbgmm_b(xa,
                         paged_lora_b[..., :out_head_size],
                         q_start_loc=start_loc,
                         q_seqlens=seq_lens,
                         adapter_ids=adapter_ids,
                         scaling=scaling,
                         rank_page_table=page_table,
                         rank_page_start=page_start,
                         ranks=ranks,
                         max_seq_len=max_seq_len,
                         max_rank=max_rank)

        torch.testing.assert_allclose(gt, output)