jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_jagged_softmax(
        self,
        B: int,
        max_L: int,
        D: int,
        dtype: torch.dtype,
        device_type: str,
    ) -> None:
        device = torch.device(device_type)
        torch.backends.cuda.matmul.allow_tf32 = False
        lengths = torch.randint(max_L + 1, size=(B,), device=device)
        total_length = int(lengths.sum().item())
        offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths)
        values = torch.rand(
            (total_length, D), requires_grad=True, dtype=dtype, device=device
        )
        output, _ = torch.ops.fbgemm.jagged_softmax(
            values,
            offsets,
            max_L,
        )
        values_ref = values.detach().clone().requires_grad_(True)
        output_ref, _ = torch.ops.fbgemm.dense_to_jagged(
            torch.nn.functional.softmax(
                torch.ops.fbgemm.jagged_to_padded_dense(
                    values_ref,
                    [offsets],
                    max_lengths=[max_L],
                    padding_value=-5e7,
                ).transpose(1, 2),
                dim=-1,
            ).permute(0, 2, 1),
            [offsets],
            total_length,
        )

        # verify forward
        torch.testing.assert_allclose(output, output_ref)

        # verify backward
        grad_output = output.detach().clone().requires_grad_(True)

        output.backward(grad_output)
        output_ref.backward(grad_output)

        torch.testing.assert_allclose(values.grad, values_ref.grad)
//  [python]
def test_jagged_softmax(
        self,
        B: int,
        max_L: int,
        D: int,
        dtype: torch.dtype,
        device_type: str,
    ) -> None:
        device = torch.device(device_type)
        torch.backends.cuda.matmul.allow_tf32 = False
        lengths = torch.randint(max_L + 1, size=(B,), device=device)
        total_length = int(lengths.sum().item())
        offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths)
        values = torch.rand(
            (total_length, D), requires_grad=True, dtype=dtype, device=device
        )
        output, _ = torch.ops.fbgemm.jagged_softmax(
            values,
            offsets,
            max_L,
        )
        values_ref = values.detach().clone().requires_grad_(True)
        output_ref, _ = torch.ops.fbgemm.dense_to_jagged(
            torch.nn.functional.softmax(
                torch.ops.fbgemm.jagged_to_padded_dense(
                    values_ref,
                    [offsets],
                    max_lengths=[max_L],
                    padding_value=-5e7,
                ).transpose(1, 2),
                dim=-1,
            ).permute(0, 2, 1),
            [offsets],
            total_length,
        )

        # verify forward
        torch.testing.assert_close(output, output_ref)

        # verify backward
        grad_output = output.detach().clone().requires_grad_(True)

        output.backward(grad_output)
        output_ref.backward(grad_output)

        torch.testing.assert_close(values.grad, values_ref.grad)
//  [python]
def test_masked_select_jagged_1d(
        self,
        max_length: int,
        batch_size: int,
        index_dtype: torch.dtype,
        jagged_tensor_dtype: torch.dtype,
        empty_lengths: bool,
        use_cpu: bool,
    ) -> None:
        device = "cpu" if use_cpu else "cuda"
        if empty_lengths:
            lengths = torch.zeros(batch_size, dtype=index_dtype, device=device)
        else:
            lengths = torch.randint(
                low=0,
                high=max_length + 1,
                size=(batch_size,),
                dtype=index_dtype,
                device=device,
            )
        lengths[batch_size // 2] = 0  # test a corner case
        n = int(lengths.sum().item())
        values = torch.randint(
            2**16,
            (n,),
            dtype=jagged_tensor_dtype,
            device=device,
        )
        mask = torch.randint(2, (n,)) > 0

        masked_values, masked_lengths = torch.ops.fbgemm.masked_select_jagged_1d(
            values,
            lengths,
            mask,
        )

        masked_values_ref = values[mask]
        cum_count = torch.cumsum(mask, 0)
        cum_count = torch.cat((cum_count, torch.tensor([0])))
        cum_length = cum_count[torch.cumsum(lengths, 0) - 1]
        cum_length_shift_right = torch.roll(cum_length, 1)
        cum_length_shift_right[0] = 0
        masked_lengths_ref = (cum_length - cum_length_shift_right).to(lengths.dtype)

        torch.testing.assert_allclose(masked_values, masked_values_ref)
        torch.testing.assert_allclose(masked_lengths, masked_lengths_ref)
//  [python]
def test_masked_select_jagged_1d(
        self,
        max_length: int,
        batch_size: int,
        index_dtype: torch.dtype,
        jagged_tensor_dtype: torch.dtype,
        empty_lengths: bool,
        use_cpu: bool,
    ) -> None:
        device = "cpu" if use_cpu else "cuda"
        if empty_lengths:
            lengths = torch.zeros(batch_size, dtype=index_dtype, device=device)
        else:
            lengths = torch.randint(
                low=0,
                high=max_length + 1,
                size=(batch_size,),
                dtype=index_dtype,
                device=device,
            )
        lengths[batch_size // 2] = 0  # test a corner case
        n = int(lengths.sum().item())
        values = torch.randint(
            2**16,
            (n,),
            dtype=jagged_tensor_dtype,
            device=device,
        )
        mask = torch.randint(2, (n,)) > 0

        masked_values, masked_lengths = torch.ops.fbgemm.masked_select_jagged_1d(
            values,
            lengths,
            mask,
        )

        masked_values_ref = values[mask]
        cum_count = torch.cumsum(mask, 0)
        cum_count = torch.cat((cum_count, torch.tensor([0])))
        cum_length = cum_count[torch.cumsum(lengths, 0) - 1]
        cum_length_shift_right = torch.roll(cum_length, 1)
        cum_length_shift_right[0] = 0
        masked_lengths_ref = (cum_length - cum_length_shift_right).to(lengths.dtype)

        torch.testing.assert_close(masked_values, masked_values_ref)
        torch.testing.assert_close(masked_lengths, masked_lengths_ref)
//  [python]
def get_multi_weights_col(self, prefixes: List[str], quantize: str, dim: int):
        if quantize in ["gptq", "awq"]:
            try:
                qweight = torch.cat(
                    [self.get_sharded(f"{p}.qweight", dim=1) for p in prefixes], dim=1
                )
            except RuntimeError:
                raise RuntimeError(
                    f"Cannot load `{quantize}` weight, make sure the model is already quantized"
                )

            qzeros = torch.cat(
                [self.get_sharded(f"{p}.qzeros", dim=1) for p in prefixes], dim=1
            )
            scales = torch.cat(
                [self.get_sharded(f"{p}.scales", dim=1) for p in prefixes], dim=1
            )

            bits, groupsize, desc_act, quant_method = self._get_gptq_params()

            from text_generation_server.utils.layers import HAS_EXLLAMA

            use_exllama = (
                bits == 4 and HAS_EXLLAMA and quantize == "gptq" and not desc_act
            )

            if quantize == "gptq" and quant_method == "gptq":
                w = [self.get_tensor(f"{p}.g_idx") for p in prefixes]
                for w2 in w[1:]:
                    torch.allclose(w2, w[0])
                g_idx = w[0]
            elif quantize == "gptq" and quant_method == "awq":
                log_once(
                    logger.info, "Converting AWQ model to Exllama/GPTQ packing format."
                )
                from text_generation_server.utils.awq.conversion_utils import (
                    fast_awq_to_gptq,
                )

                qweight, qzeros = fast_awq_to_gptq(qweight, qzeros)
                if use_exllama:
                    g_idx = None
                else:
                    g_idx = (
                        torch.arange(
                            qweight.shape[0] * (32 // bits), device=qweight.device
                        )
                        // groupsize
                    ).to(dtype=torch.int32)
            else:
                g_idx = None

            weight = (qweight, qzeros, scales, g_idx, bits, groupsize, use_exllama)
        else:
            w = [self.get_sharded(f"{p}.weight", dim=0) for p in prefixes]
            weight = torch.cat(w, dim=dim)
        return weight
//  [python]
def get_multi_weights_col(self, prefixes: List[str], quantize: str, dim: int):
        if quantize in ["gptq", "awq"]:
            try:
                qweight = torch.cat(
                    [self.get_sharded(f"{p}.qweight", dim=1) for p in prefixes], dim=1
                )
            except RuntimeError:
                raise RuntimeError(
                    f"Cannot load `{quantize}` weight, make sure the model is already quantized"
                )

            qzeros = torch.cat(
                [self.get_sharded(f"{p}.qzeros", dim=1) for p in prefixes], dim=1
            )
            scales = torch.cat(
                [self.get_sharded(f"{p}.scales", dim=1) for p in prefixes], dim=1
            )

            bits, groupsize, desc_act, quant_method = self._get_gptq_params()

            from text_generation_server.utils.layers import HAS_EXLLAMA

            use_exllama = (
                bits == 4 and HAS_EXLLAMA and quantize == "gptq" and not desc_act
            )

            if quantize == "gptq" and quant_method == "gptq":
                w = [self.get_tensor(f"{p}.g_idx") for p in prefixes]
                for w2 in w[1:]:
                    torch.testing.assert_close(w2, w[0])
                g_idx = w[0]
            elif quantize == "gptq" and quant_method == "awq":
                log_once(
                    logger.info, "Converting AWQ model to Exllama/GPTQ packing format."
                )
                from text_generation_server.utils.awq.conversion_utils import (
                    fast_awq_to_gptq,
                )

                qweight, qzeros = fast_awq_to_gptq(qweight, qzeros)
                if use_exllama:
                    g_idx = None
                else:
                    g_idx = (
                        torch.arange(
                            qweight.shape[0] * (32 // bits), device=qweight.device
                        )
                        // groupsize
                    ).to(dtype=torch.int32)
            else:
                g_idx = None

            weight = (qweight, qzeros, scales, g_idx, bits, groupsize, use_exllama)
        else:
            w = [self.get_sharded(f"{p}.weight", dim=0) for p in prefixes]
            weight = torch.cat(w, dim=dim)
        return weight
//  [python]
def test_save_and_load_state_dict(self):
        class TestModule(torch.nn.Module):
            def forward(self, x):
                return x + x

        inputs = [torch.randn(1, 1)]
        mod = TestModule().eval()
        ref_output = mod(*inputs)

        mod = acc_tracer.trace(mod, inputs)
        interp = TRTInterpreter(mod, input_specs=InputTensorSpec.from_tensors(inputs))
        trt_mod = TRTModule(*interp.run(lower_precision=LowerPrecision.FP32))
        st = trt_mod.state_dict()

        new_trt_mod = TRTModule()
        new_trt_mod.load_state_dict(st)

        torch.testing.assert_allclose(
            new_trt_mod(inputs[0].cuda()).cpu(), ref_output, rtol=1e-04, atol=1e-04
        )
//  [python]
def test_save_and_load_state_dict(self):
        class TestModule(torch.nn.Module):
            def forward(self, x):
                return x + x

        inputs = [torch.randn(1, 1)]
        mod = TestModule().eval()
        ref_output = mod(*inputs)

        mod = acc_tracer.trace(mod, inputs)
        interp = TRTInterpreter(mod, input_specs=InputTensorSpec.from_tensors(inputs))
        trt_mod = TRTModule(*interp.run(lower_precision=LowerPrecision.FP32))
        st = trt_mod.state_dict()

        new_trt_mod = TRTModule()
        new_trt_mod.load_state_dict(st)

        torch.testing.assert_close(
            new_trt_mod(inputs[0].cuda()).cpu(), ref_output, rtol=1e-04, atol=1e-04
        )
//  [python]
def check_func(func1, func2, args=(), name="", niters=100, rtol=None, atol=None):
    """Verify that both input functions produce identical outputs

    Args:
        func1: First function to test
        func2: Second function to test
        args: Arguments to both functions
        name: Name of this test
        niters: Number of test iterations (default 5)
        rtol: Relative tolerance (by default, selected based on datatype)
        atol: Absolute tolerance (by default, selected based on datatype)
    """
    # Make sure cuda is already loaded
    torch.zeros(1, dtype=torch.float32, device="cuda")

    all_t1 = []
    all_t2 = []
    for i in range(niters):
        torch.cuda.synchronize()
        t1 = time.time()
        out1 = func1(*args)
        torch.cuda.synchronize()
        all_t1.append(time.time() - t1)

        torch.cuda.synchronize()
        t2 = time.time()
        out2 = func2(*args)
        torch.cuda.synchronize()
        all_t2.append(time.time() - t2)

        try:
            assert type(out1) == type(out2)
            if isinstance(out1, torch.Tensor) and isinstance(out2, torch.Tensor):
                torch.allclose(out1, out2, rtol=rtol, atol=atol)
            elif isinstance(out1, np.ndarray) and isinstance(out2, np.ndarray):
                np.testing.assert_allclose(out1, out2, rtol=rtol, atol=atol)
            else:
                assert all(
                    elt1 == elt2 for elt1, elt2 in zip(out1, out2)
                ), f"out1={out1} but out2={out2}"
        except Exception as e:
            print(f"Error: {e}")

    all_t1 = all_t1[10:]  # Remove the first few iterations to account for warmup
    all_t2 = all_t2[10:]
    avg_t1 = 1000 * mean(all_t1)  # milliseconds
    avg_t2 = 1000 * mean(all_t2)
    std_t1 = 1000 * stdev(all_t1) if len(all_t1) > 1 else 0
    std_t2 = 1000 * stdev(all_t2) if len(all_t1) > 1 else 0

    print(
        f"Test '{name}' passed:\tavg_t1={avg_t1:.2f}ms,\tavg_t2={avg_t2:.2f}ms,"
        f"\tstd_t1={std_t1:.2f}ms,\tstd_t2={std_t2:.2f}ms"
    )
//  [python]
def check_func(func1, func2, args=(), name="", niters=100, rtol=None, atol=None):
    """Verify that both input functions produce identical outputs

    Args:
        func1: First function to test
        func2: Second function to test
        args: Arguments to both functions
        name: Name of this test
        niters: Number of test iterations (default 5)
        rtol: Relative tolerance (by default, selected based on datatype)
        atol: Absolute tolerance (by default, selected based on datatype)
    """
    # Make sure cuda is already loaded
    torch.zeros(1, dtype=torch.float32, device="cuda")

    all_t1 = []
    all_t2 = []
    for i in range(niters):
        torch.cuda.synchronize()
        t1 = time.time()
        out1 = func1(*args)
        torch.cuda.synchronize()
        all_t1.append(time.time() - t1)

        torch.cuda.synchronize()
        t2 = time.time()
        out2 = func2(*args)
        torch.cuda.synchronize()
        all_t2.append(time.time() - t2)

        try:
            assert type(out1) == type(out2)
            if isinstance(out1, torch.Tensor) and isinstance(out2, torch.Tensor):
                torch.testing.assert_close(out1, out2, rtol=rtol, atol=atol)
            elif isinstance(out1, np.ndarray) and isinstance(out2, np.ndarray):
                np.testing.assert_allclose(out1, out2, rtol=rtol, atol=atol)
            else:
                assert all(
                    elt1 == elt2 for elt1, elt2 in zip(out1, out2)
                ), f"out1={out1} but out2={out2}"
        except Exception as e:
            print(f"Error: {e}")

    all_t1 = all_t1[10:]  # Remove the first few iterations to account for warmup
    all_t2 = all_t2[10:]
    avg_t1 = 1000 * mean(all_t1)  # milliseconds
    avg_t2 = 1000 * mean(all_t2)
    std_t1 = 1000 * stdev(all_t1) if len(all_t1) > 1 else 0
    std_t2 = 1000 * stdev(all_t2) if len(all_t1) > 1 else 0

    print(
        f"Test '{name}' passed:\tavg_t1={avg_t1:.2f}ms,\tavg_t2={avg_t2:.2f}ms,"
        f"\tstd_t1={std_t1:.2f}ms,\tstd_t2={std_t2:.2f}ms"
    )
//  [python]
def test_from_filename(self, shape, dtype, device, tmp_path, from_path):
        if from_path:
            filename = tmp_path / "file.memmap"
        else:
            filename = None
        if dtype is None:
            dtype = torch.float32
        tensor = -torch.randint(10, shape, dtype=dtype, device=device)
        t = MemoryMappedTensor.full_like(tensor, 2, filename=filename)
        if filename is not None:
            t2 = MemoryMappedTensor.from_filename(filename, dtype=dtype, shape=shape)
        else:
            t2 = MemoryMappedTensor.from_handler(
                t._handler, dtype=dtype, shape=shape, index=None
            )
        torch.testing.assert_allclose(t, t2)
//  [python]
def test_from_filename(self, shape, dtype, device, tmp_path, from_path):
        if from_path:
            filename = tmp_path / "file.memmap"
        else:
            filename = None
        if dtype is None:
            dtype = torch.float32
        tensor = -torch.randint(10, shape, dtype=dtype, device=device)
        t = MemoryMappedTensor.full_like(tensor, 2, filename=filename)
        if filename is not None:
            t2 = MemoryMappedTensor.from_filename(filename, dtype=dtype, shape=shape)
        else:
            t2 = MemoryMappedTensor.from_handler(
                t._handler, dtype=dtype, shape=shape, index=None
            )
        torch.testing.assert_close(t, t2)
//  [python]
def test_stack_tensors_3d(self):
        tensors = [torch.ones([2, 2]), torch.ones([1, 3])]
        max_lens = [4, 2]
        expected_output = torch.tensor(
            [[[1, 1, 0, 0], [1, 1, 0, 0]], [[1, 1, 1, 0], [0, 0, 0, 0]]], dtype=torch.float32
        )

        stacked_tensor = stack_tensors(tensors=tensors, max_lens=max_lens)

        self.assertEqual(stacked_tensor.tolist(), expected_output.tolist())
//  [python]
def test_stack_tensors_3d(self):
        tensors = [torch.ones([2, 2]), torch.ones([1, 3])]
        max_lens = [4, 2]
        expected_output = torch.tensor(
            [[[1, 1, 0, 0], [1, 1, 0, 0]], [[1, 1, 1, 0], [0, 0, 0, 0]]], dtype=torch.float32
        )

        stacked_tensor = stack_tensors(tensors=tensors, max_lens=max_lens)

        torch.testing.assert_close(stacked_tensor, expected_output)
//  [python]
def test_merge_attentions_against_ref(bmghk: bool, stack_inputs: bool):
    split_k = 16
    B = 12
    M = 137
    G = 2 if bmghk else 1
    N_H_L = 8
    D_H = 128
    dtype = torch.float32

    attn_split = torch.randn([split_k, B, M, G, N_H_L, D_H], dtype=dtype, device="cuda")
    lse_split = torch.randn([split_k, B, G, N_H_L, M], dtype=dtype, device="cuda")

    if not bmghk:
        attn_split = attn_split[:, :, :, 0]
        lse_split = lse_split[:, :, 0]
    attn_out_ref, lse_out_ref = _merge_attentions_ref(attn_split, lse_split)
    attn_split_ = attn_split if stack_inputs else attn_split.unbind(0)
    lse_split_ = lse_split if stack_inputs else lse_split.unbind(0)
    attn_out, lse_out = fmha.merge_attentions(attn_split_, lse_split_)  # type: ignore

    torch.testing.assert_allclose(lse_out, lse_out_ref, rtol=1e-4, atol=1e-4)
    torch.testing.assert_allclose(attn_out, attn_out_ref, rtol=1e-4, atol=1e-4)
//  [python]
def test_merge_attentions_against_ref(bmghk: bool, stack_inputs: bool):
    split_k = 16
    B = 12
    M = 137
    G = 2 if bmghk else 1
    N_H_L = 8
    D_H = 128
    dtype = torch.float32

    attn_split = torch.randn([split_k, B, M, G, N_H_L, D_H], dtype=dtype, device="cuda")
    lse_split = torch.randn([split_k, B, G, N_H_L, M], dtype=dtype, device="cuda")

    if not bmghk:
        attn_split = attn_split[:, :, :, 0]
        lse_split = lse_split[:, :, 0]
    attn_out_ref, lse_out_ref = _merge_attentions_ref(attn_split, lse_split)
    attn_split_ = attn_split if stack_inputs else attn_split.unbind(0)
    lse_split_ = lse_split if stack_inputs else lse_split.unbind(0)
    attn_out, lse_out = fmha.merge_attentions(attn_split_, lse_split_)  # type: ignore

    torch.testing.assert_close(lse_out, lse_out_ref, rtol=1e-4, atol=1e-4)
    torch.testing.assert_close(attn_out, attn_out_ref, rtol=1e-4, atol=1e-4)
//  [python]
def test_stack_tensors(self):
        tensors = [torch.ones([2]), torch.ones([4]), torch.ones([3])]
        max_lens = [6]
        expected_output = torch.tensor(
            [[1, 1, 0, 0, 0, 0], [1, 1, 1, 1, 0, 0], [1, 1, 1, 0, 0, 0]], dtype=torch.float32
        )

        stacked_tensor = stack_tensors(tensors=tensors, max_lens=max_lens)

        self.assertEqual(stacked_tensor.tolist(), expected_output.tolist())
//  [python]
def test_stack_tensors(self):
        tensors = [torch.ones([2]), torch.ones([4]), torch.ones([3])]
        max_lens = [6]
        expected_output = torch.tensor(
            [[1, 1, 0, 0, 0, 0], [1, 1, 1, 1, 0, 0], [1, 1, 1, 0, 0, 0]], dtype=torch.float32
        )

        stacked_tensor = stack_tensors(tensors=tensors, max_lens=max_lens)

        torch.testing.assert_close(stacked_tensor, expected_output)
//  [python]
def test_sine_axis(batch_size):
    set_global_params({"_faster_log_maps": True})
    if not torch.cuda.is_available():
        return
    for _ in range(10):
        g = so3_impl.rand(*batch_size)
        g_cuda = g.to("cuda:0")
        sa_1 = so3_impl._sine_axis_fn(g, g.shape[:-2])
        sa_2 = so3_impl._sine_axis_fn(g_cuda, g.shape[:-2])
        torch.testing.assert_allclose(sa_1, sa_2.cpu())
//  [python]
def test_sine_axis(batch_size):
    set_global_params({"_faster_log_maps": True})
    if not torch.cuda.is_available():
        return
    for _ in range(10):
        g = so3_impl.rand(*batch_size)
        g_cuda = g.to("cuda:0")
        sa_1 = so3_impl._sine_axis_fn(g, g.shape[:-2])
        sa_2 = so3_impl._sine_axis_fn(g_cuda, g.shape[:-2])
        torch.testing.assert_close(sa_1, sa_2.cpu())
//  [python]
def test_vee(batch_size: Union[int, Sequence[int]], dtype: torch.dtype):
    if isinstance(batch_size, int):
        batch_size = (batch_size,)
    rng = torch.Generator()
    rng.manual_seed(0)
    tangent_vector = torch.rand(*batch_size, 3, dtype=dtype, generator=rng)
    matrix = so3_impl._hat_autograd_fn(tangent_vector)

    # check analytic backward for the operator
    check_lie_group_function(so3_impl, "vee", TEST_EPS, (matrix,))

    # check the correctness of hat and vee
    actual_tangent_vector = so3_impl._vee_autograd_fn(matrix)
    t.assert_close(
        actual_tangent_vector, tangent_vector, atol=TEST_EPS, rtol=TEST_EPS
    )
//  [python]
def test_vee(batch_size: Union[int, Sequence[int]], dtype: torch.dtype):
    if isinstance(batch_size, int):
        batch_size = (batch_size,)
    rng = torch.Generator()
    rng.manual_seed(0)
    tangent_vector = torch.rand(*batch_size, 3, dtype=dtype, generator=rng)
    matrix = so3_impl._hat_autograd_fn(tangent_vector)

    # check analytic backward for the operator
    check_lie_group_function(so3_impl, "vee", TEST_EPS, (matrix,))

    # check the correctness of hat and vee
    actual_tangent_vector = so3_impl._vee_autograd_fn(matrix)
    torch.testing.assert_close(
        actual_tangent_vector, tangent_vector, atol=TEST_EPS, rtol=TEST_EPS
    )
//  [python]
def _assert_ghost_batches_normalized(
        self,
        module: ModuleWithBatchnorm,
        ghost_batch_size: int,
        batch_size: int,
    ) -> None:
        torch.manual_seed(123)
        size = [batch_size, module.num_features] + ([3] * module.num_dims)
        X = torch.randn(size=size)
        module.train()
        out = module(X)
        n_channels = out.shape[1]
        # reduce over everything but channel idx
        reduce_dims = (0,) + tuple(range(2, out.ndim))

        nchunks = int(math.ceil(batch_size / ghost_batch_size))
        for ghost_batch in out.chunk(nchunks):
            channel_variances, channel_means = torch.var_mean(ghost_batch, dim=reduce_dims, unbiased=False)
            torch.testing.assert_allclose(channel_variances, torch.ones(n_channels))
            torch.testing.assert_allclose(channel_means, torch.zeros(n_channels))
//  [python]
def _assert_ghost_batches_normalized(
        self,
        module: ModuleWithBatchnorm,
        ghost_batch_size: int,
        batch_size: int,
    ) -> None:
        torch.manual_seed(123)
        size = [batch_size, module.num_features] + ([3] * module.num_dims)
        X = torch.randn(size=size)
        module.train()
        out = module(X)
        n_channels = out.shape[1]
        # reduce over everything but channel idx
        reduce_dims = (0,) + tuple(range(2, out.ndim))

        nchunks = int(math.ceil(batch_size / ghost_batch_size))
        for ghost_batch in out.chunk(nchunks):
            channel_variances, channel_means = torch.var_mean(ghost_batch, dim=reduce_dims, unbiased=False)
            torch.testing.assert_close(channel_variances, torch.ones(n_channels))
            torch.testing.assert_close(channel_means, torch.zeros(n_channels))
//  [python]
def test_auroc_base(self) -> None:
        num_classes = 4
        input = 10 * torch.randn(BATCH_SIZE, num_classes)
        input_prob = input.abs() / input.abs().sum(dim=-1, keepdim=True)
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        compute_result = torch.tensor(
            roc_auc_score(target, input_prob, average="macro", multi_class="ovr"),
            dtype=torch.float32,
        )
        my_compute_result = multiclass_auroc(
            input_prob, target, num_classes=num_classes
        )
        t.assert_allclose(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_auroc_base(self) -> None:
        num_classes = 4
        input = 10 * torch.randn(BATCH_SIZE, num_classes)
        input_prob = input.abs() / input.abs().sum(dim=-1, keepdim=True)
        target = torch.randint(high=num_classes, size=(BATCH_SIZE,))
        compute_result = torch.tensor(
            roc_auc_score(target, input_prob, average="macro", multi_class="ovr"),
            dtype=torch.float32,
        )
        my_compute_result = multiclass_auroc(
            input_prob, target, num_classes=num_classes
        )
        torch.testing.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_distributed_collectives_fn(strategy, collective):
    collective.create_group()

    # all_gather
    tensor_list = [torch.zeros(2, dtype=torch.long) for _ in range(strategy.num_processes)]
    this = torch.arange(2, dtype=torch.long) + 2 * strategy.global_rank
    out = collective.all_gather(tensor_list, this)
    expected = torch.arange(2 * strategy.num_processes).split(2)
    t.assert_close(tuple(out), expected)

    # reduce
    this = torch.tensor(strategy.global_rank + 1)
    out = collective.reduce(this, dst=0, op="max")
    expected = torch.tensor(strategy.num_processes) if strategy.global_rank == 0 else this
    t.assert_close(out, expected)

    # all_reduce
    this = torch.tensor(strategy.global_rank + 1)
    out = collective.all_reduce(this, op=ReduceOp.MIN)
    expected = torch.tensor(1)
    t.assert_close(out, expected)
//  [python]
def _test_distributed_collectives_fn(strategy, collective):
    collective.create_group()

    # all_gather
    tensor_list = [torch.zeros(2, dtype=torch.long) for _ in range(strategy.num_processes)]
    this = torch.arange(2, dtype=torch.long) + 2 * strategy.global_rank
    out = collective.all_gather(tensor_list, this)
    expected = torch.arange(2 * strategy.num_processes).split(2)
    torch.testing.assert_close(tuple(out), expected)

    # reduce
    this = torch.tensor(strategy.global_rank + 1)
    out = collective.reduce(this, dst=0, op="max")
    expected = torch.tensor(strategy.num_processes) if strategy.global_rank == 0 else this
    torch.testing.assert_close(out, expected)

    # all_reduce
    this = torch.tensor(strategy.global_rank + 1)
    out = collective.all_reduce(this, op=ReduceOp.MIN)
    expected = torch.tensor(1)
    torch.testing.assert_close(out, expected)
//  [python]
def _test_auroc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_tasks: int = 1,
        weight: Optional[torch.Tensor] = None,
        compute_result: Optional[torch.Tensor] = None,
        use_fbgemm: Optional[bool] = False,
    ) -> None:
        if compute_result is None:
            compute_result = (
                torch.tensor(roc_auc_score(target, input))
                if weight is None
                else torch.tensor(roc_auc_score(target, input, sample_weight=weight))
            )
        if torch.cuda.is_available():
            my_compute_result = binary_auroc(
                input.to(device="cuda"),
                target.to(device="cuda"),
                weight=weight if weight is None else weight.to(device="cuda"),
                num_tasks=num_tasks,
                use_fbgemm=use_fbgemm,
            )
            compute_result = compute_result.to(device="cuda")
        else:
            my_compute_result = binary_auroc(input, target, num_tasks=num_tasks)

        t.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_auroc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_tasks: int = 1,
        weight: Optional[torch.Tensor] = None,
        compute_result: Optional[torch.Tensor] = None,
        use_fbgemm: Optional[bool] = False,
    ) -> None:
        if compute_result is None:
            compute_result = (
                torch.tensor(roc_auc_score(target, input))
                if weight is None
                else torch.tensor(roc_auc_score(target, input, sample_weight=weight))
            )
        if torch.cuda.is_available():
            my_compute_result = binary_auroc(
                input.to(device="cuda"),
                target.to(device="cuda"),
                weight=weight if weight is None else weight.to(device="cuda"),
                num_tasks=num_tasks,
                use_fbgemm=use_fbgemm,
            )
            compute_result = compute_result.to(device="cuda")
        else:
            my_compute_result = binary_auroc(input, target, num_tasks=num_tasks)

        torch.testing.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_conversions_all_patterns(self, device, dtype):
        r, c = 32, 128

        dense_inv, dense_val = rand_sparse_semi_structured_all_patterns(r, c, dtype, device)

        compressed = to_sparse_semi_structured(dense_inv)
        dense = compressed.to_dense()

        torch.testing.assert_allclose(dense, dense_val, rtol=0, atol=0)
//  [python]
def test_conversions_all_patterns(self, device, dtype):
        r, c = 32, 128

        dense_inv, dense_val = rand_sparse_semi_structured_all_patterns(r, c, dtype, device)

        compressed = to_sparse_semi_structured(dense_inv)
        dense = compressed.to_dense()

        torch.testing.assert_close(dense, dense_val, rtol=0, atol=0)
//  [python]
def compare_lora(
    bsz: int,
    seq_len: int,
    in_dim: int,
    out_dim: int,
    rank: int,
    alpha: float,
    dropout: float,
) -> None:

    # make sure we have the right seed for generating outputs
    # this should match up the seed value set in the corresponding
    # unit test
    torch.manual_seed(16)

    # generate input tensor used by both implementations
    x_input = torch.randn(bsz, seq_len, in_dim)

    # Initialize our implementation
    lora = LoRALinear(
        in_dim=in_dim,
        out_dim=out_dim,
        rank=rank,
        alpha=alpha,
        use_bias=True,
        dropout=dropout,
    )
    fixed_init_model(lora)

    with torch.no_grad():
        output = lora(x_input)

    # Initialize reference implementation
    lora_ref = LoRALinearRef(
        in_features=in_dim,
        out_features=out_dim,
        r=rank,
        lora_alpha=alpha,
        lora_dropout=dropout,
    )

    sd_mapping = {
        "weight": "weight",
        "bias": "bias",
        "lora_a.weight": "lora_A",
        "lora_b.weight": "lora_B",
    }
    mapped_sd = {sd_mapping.get(k): v for k, v in lora.state_dict().items()}
    lora_ref.load_state_dict(mapped_sd)
    with torch.no_grad():
        output_ref = lora_ref(x_input)

    print(output_ref.mean())
    torch.testing.assert_allclose(output_ref, output, atol=1e-6, rtol=1e-6)
//  [python]
def compare_lora(
    bsz: int,
    seq_len: int,
    in_dim: int,
    out_dim: int,
    rank: int,
    alpha: float,
    dropout: float,
) -> None:

    # make sure we have the right seed for generating outputs
    # this should match up the seed value set in the corresponding
    # unit test
    torch.manual_seed(16)

    # generate input tensor used by both implementations
    x_input = torch.randn(bsz, seq_len, in_dim)

    # Initialize our implementation
    lora = LoRALinear(
        in_dim=in_dim,
        out_dim=out_dim,
        rank=rank,
        alpha=alpha,
        use_bias=True,
        dropout=dropout,
    )
    fixed_init_model(lora)

    with torch.no_grad():
        output = lora(x_input)

    # Initialize reference implementation
    lora_ref = LoRALinearRef(
        in_features=in_dim,
        out_features=out_dim,
        r=rank,
        lora_alpha=alpha,
        lora_dropout=dropout,
    )

    sd_mapping = {
        "weight": "weight",
        "bias": "bias",
        "lora_a.weight": "lora_A",
        "lora_b.weight": "lora_B",
    }
    mapped_sd = {sd_mapping.get(k): v for k, v in lora.state_dict().items()}
    lora_ref.load_state_dict(mapped_sd)
    with torch.no_grad():
        output_ref = lora_ref(x_input)

    print(output_ref.mean())
    torch.testing.assert_close(output_ref, output, atol=1e-6, rtol=1e-6)
//  [python]
def test_auroc_average_options(self) -> None:
        input = torch.tensor(
            [
                [0.16, 0.04, 0.8],
                [0.1, 0.7, 0.2],
                [0.16, 0.8, 0.04],
                [0.16, 0.04, 0.8],
            ]
        )
        target = torch.tensor([0, 0, 1, 2])
        # average = macro
        compute_result = torch.tensor(
            roc_auc_score(target, input, average="macro", multi_class="ovr"),
            dtype=torch.float32,
        )
        my_compute_result = multiclass_auroc(
            input, target, num_classes=3, average="macro"
        )
        torch.testing.assert_allclose(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # average = None
        # sklearn.metrics.roc_auc_score does not support average=None
        expected_compute_result = torch.tensor([0.2500, 1.0000, 5 / 6])
        my_compute_result = multiclass_auroc(input, target, num_classes=3, average=None)
        torch.testing.assert_allclose(my_compute_result, expected_compute_result)
//  [python]
def test_auroc_average_options(self) -> None:
        input = torch.tensor(
            [
                [0.16, 0.04, 0.8],
                [0.1, 0.7, 0.2],
                [0.16, 0.8, 0.04],
                [0.16, 0.04, 0.8],
            ]
        )
        target = torch.tensor([0, 0, 1, 2])
        # average = macro
        compute_result = torch.tensor(
            roc_auc_score(target, input, average="macro", multi_class="ovr"),
            dtype=torch.float32,
        )
        my_compute_result = multiclass_auroc(
            input, target, num_classes=3, average="macro"
        )
        torch.testing.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )

        # average = None
        # sklearn.metrics.roc_auc_score does not support average=None
        expected_compute_result = torch.tensor([0.2500, 1.0000, 5 / 6])
        my_compute_result = multiclass_auroc(input, target, num_classes=3, average=None)
        torch.testing.assert_close(my_compute_result, expected_compute_result)
//  [python]
def test_conversions(self, device, dtype):

        def run_test(r, c, device, dtype):
            dense_ref = rand_sparse_semi_structured(r, c, dtype, device)

            compressed = to_sparse_semi_structured(dense_ref)

            # The torch.ops.aten._to_sparse_semi_structured operator
            # uses CUTLASS to perform conversion from given dense
            # matrix to the pair of corresponding sparse and metadata
            # matrices, with the later used here as a reference to
            # compare the metadata matrix produced by conversion
            # performed by SparseSemiStructuredTensor class
            # constructor against.
            _, meta_ref = torch.ops.aten._to_sparse_semi_structured(dense_ref)

            meta = compressed.indices()
            t.assert_close(meta, meta_ref, rtol=0, atol=0)

            dense = compressed.to_dense()
            t.assert_close(dense, dense_ref, rtol=0, atol=0)

        shapes = [[32, 128], [32, 256], [64, 128], [64, 256]]
        for r, c in shapes:
            run_test(r, c, device, dtype)
//  [python]
def test_conversions(self, device, dtype):

        def run_test(r, c, device, dtype):
            dense_ref = rand_sparse_semi_structured(r, c, dtype, device)

            compressed = to_sparse_semi_structured(dense_ref)

            # The torch.ops.aten._to_sparse_semi_structured operator
            # uses CUTLASS to perform conversion from given dense
            # matrix to the pair of corresponding sparse and metadata
            # matrices, with the later used here as a reference to
            # compare the metadata matrix produced by conversion
            # performed by SparseSemiStructuredTensor class
            # constructor against.
            _, meta_ref = torch.ops.aten._to_sparse_semi_structured(dense_ref)

            meta = compressed.indices()
            torch.testing.assert_close(meta, meta_ref, rtol=0, atol=0)

            dense = compressed.to_dense()
            torch.testing.assert_close(dense, dense_ref, rtol=0, atol=0)

        shapes = [[32, 128], [32, 256], [64, 128], [64, 256]]
        for r, c in shapes:
            run_test(r, c, device, dtype)
//  [python]
def test_position_encoding(self):
        batch_size = 2
        max_time = 3
        emb_size = hidden_size = 12

        x = torch.zeros([batch_size, max_time, emb_size])
        pe = PositionalEncoding(emb_size)
        output = pe(x)
        self.assertEqual(pe.pe.size(2), hidden_size)
        for i in range(batch_size):
            self.assertTrue(torch.allclose(output[i].unsqueeze(0), pe.pe[:, :max_time], rtol=1e-4, atol=1e-4))
//  [python]
def test_position_encoding(self):
        batch_size = 2
        max_time = 3
        emb_size = hidden_size = 12

        x = torch.zeros([batch_size, max_time, emb_size])
        pe = PositionalEncoding(emb_size)
        output = pe(x)
        self.assertEqual(pe.pe.size(2), hidden_size)
        for i in range(batch_size):
            torch.testing.assert_close(
                output[i].unsqueeze(0), pe.pe[:, :max_time], rtol=1e-4, atol=1e-4
            )
//  [python]
def test_fx2ait_lower(self):
        class TestMod(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, x):
                values = torch.sigmoid(x)
                return get_length(values)

        mod = TestMod().half().cuda()
        input = [torch.randn(2, 3).half().cuda()]
        ref_output = mod(*input)
        lower = AitLowerer.create(
            LowerSettings(workdir="/tmp", name="test_ait_lower", min_acc_module_size=0)
        )
        lowered = lower(mod, input)
        lower_output = lowered(*input)
        self.assertTrue(len(lowered._modules.keys()), 2)
        self.assertEqual(ref_output, lower_output)
//  [python]
def test_fx2ait_lower(self):
        class TestMod(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, x):
                values = torch.sigmoid(x)
                return get_length(values)

        mod = TestMod().half().cuda()
        input = [torch.randn(2, 3).half().cuda()]
        ref_output = mod(*input)
        lower = AitLowerer.create(
            LowerSettings(workdir="/tmp", name="test_ait_lower", min_acc_module_size=0)
        )
        lowered = lower(mod, input)
        lower_output = lowered(*input)
        self.assertTrue(len(lowered._modules.keys()), 2)
        torch.testing.assert_close(ref_output, lower_output, check_dtype=False)
//  [python]
def test_fx2ait_lower_avoids_copies(self):
        class TestMod(torch.nn.Module):
            def forward(self, x):
                a = unsupported_op(x)
                b = a.unsqueeze(0)
                return unsupported_op(b)

        mod = TestMod().half().cuda()
        x = torch.randn((1,)).half().cuda()
        ref_output = mod(x)
        lowerer = AitLowerer.create(
            LowerSettings(
                workdir="/tmp",
                name="test_ait_lower_avoids_copies",
                min_acc_module_size=0,
            )
        )
        lowered = lowerer(mod, [x])
        lower_output = lowered(x)
        self.assertEqual(ref_output, lower_output)

        children = list(lowered.named_children())
        self.assertEqual(len(children), 1)
        name, _ = children[0]
        self.assertNotIn("_run_on_acc", name)
//  [python]
def test_fx2ait_lower_avoids_copies(self):
        class TestMod(torch.nn.Module):
            def forward(self, x):
                a = unsupported_op(x)
                b = a.unsqueeze(0)
                return unsupported_op(b)

        mod = TestMod().half().cuda()
        x = torch.randn((1,)).half().cuda()
        ref_output = mod(x)
        lowerer = AitLowerer.create(
            LowerSettings(
                workdir="/tmp",
                name="test_ait_lower_avoids_copies",
                min_acc_module_size=0,
            )
        )
        lowered = lowerer(mod, [x])
        lower_output = lowered(x)
        torch.testing.assert_close(ref_output, lower_output, check_dtype=False)

        children = list(lowered.named_children())
        self.assertEqual(len(children), 1)
        name, _ = children[0]
        self.assertNotIn("_run_on_acc", name)
//  [python]
def test(make_args, nnc=nnc_add, aten=torch.add):
    def setup(n):
        args = make_args(n)
        result_aten = aten(*args)
        result_nnc = nnc(*args)
        assert result_nnc.dtype == result_aten.dtype
        assert result_nnc.size() == result_aten.size()
        assert result_nnc.stride() == result_aten.stride()
        t.assert_close(result_aten, result_nnc)
        return (lambda: nnc(*args), lambda: aten(*args))

    return benchmark_loop(setup)
//  [python]
def test(make_args, nnc=nnc_add, aten=torch.add):
    def setup(n):
        args = make_args(n)
        result_aten = aten(*args)
        result_nnc = nnc(*args)
        assert result_nnc.dtype == result_aten.dtype
        assert result_nnc.size() == result_aten.size()
        assert result_nnc.stride() == result_aten.stride()
        torch.testing.assert_close(result_aten, result_nnc)
        return (lambda: nnc(*args), lambda: aten(*args))

    return benchmark_loop(setup)
//  [python]
def _compare_pytorch_onnx_with_ort(
    onnx_program: torch.onnx.ONNXProgram,
    model: _ModelType,
    input_args: Sequence[_InputArgsType],
    input_kwargs: Mapping[str, _InputArgsType],
    atol: Optional[float] = None,
    rtol: Optional[float] = None,
    has_mutation: bool = False,
):
    if has_mutation:
        ref_model = _try_clone_model(model)
        ref_input_args, ref_input_kwargs = _try_clone_inputs(input_args, input_kwargs)
    else:
        ref_model = model
        ref_input_args = input_args
        ref_input_kwargs = input_kwargs

    # NOTE: ONNXProgram holds a reference (not copy) to the original ref_model, including its state_dict.
    # Thus, ONNXProgram() must run before ref_model() to prevent ref_model.forward() from changing the state_dict.
    # Otherwise, the ref_model can change buffers on state_dict which would be used by ONNXProgram.__call__()
    # NOTE: `model_with_state_dict=ref_model` is specified to cover runs with FakeTensor support
    ort_outputs = onnx_program(*input_args, **input_kwargs)
    ref_outputs = ref_model(*ref_input_args, **ref_input_kwargs)
    ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(ref_outputs)

    if len(ref_outputs) != len(ort_outputs):
        raise AssertionError(
            f"Expected {len(ref_outputs)} outputs, got {len(ort_outputs)}"
        )

    for ref_output, ort_output in zip(ref_outputs, ort_outputs):
        torch.testing.assert_allclose(
            ref_output, torch.tensor(ort_output), rtol=rtol, atol=atol
        )
//  [python]
def _compare_pytorch_onnx_with_ort(
    onnx_program: torch.onnx.ONNXProgram,
    model: _ModelType,
    input_args: Sequence[_InputArgsType],
    input_kwargs: Mapping[str, _InputArgsType],
    atol: Optional[float] = None,
    rtol: Optional[float] = None,
    has_mutation: bool = False,
):
    if has_mutation:
        ref_model = _try_clone_model(model)
        ref_input_args, ref_input_kwargs = _try_clone_inputs(input_args, input_kwargs)
    else:
        ref_model = model
        ref_input_args = input_args
        ref_input_kwargs = input_kwargs

    # NOTE: ONNXProgram holds a reference (not copy) to the original ref_model, including its state_dict.
    # Thus, ONNXProgram() must run before ref_model() to prevent ref_model.forward() from changing the state_dict.
    # Otherwise, the ref_model can change buffers on state_dict which would be used by ONNXProgram.__call__()
    # NOTE: `model_with_state_dict=ref_model` is specified to cover runs with FakeTensor support
    ort_outputs = onnx_program(*input_args, **input_kwargs)
    ref_outputs = ref_model(*ref_input_args, **ref_input_kwargs)
    ref_outputs = onnx_program.adapt_torch_outputs_to_onnx(ref_outputs)

    if len(ref_outputs) != len(ort_outputs):
        raise AssertionError(
            f"Expected {len(ref_outputs)} outputs, got {len(ort_outputs)}"
        )

    for ref_output, ort_output in zip(ref_outputs, ort_outputs):
        torch.testing.assert_close(
            ref_output, torch.tensor(ort_output), rtol=rtol, atol=atol
        )