jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _assert_state_dict_equals(
        self, dict1: Dict[str, Any], dict2: Dict[str, Any]
    ) -> None:
        self.assertEqual(dict1["param_groups"], dict2["param_groups"])
        self.assertEqual(
            dict1["state"]["param_2"],
            dict2["state"]["param_2"],
        )
        self.assertEqual(
            dict1["state"]["param_1"]["tensor"].tolist(),
            dict2["state"]["param_1"]["tensor"].tolist(),
        )
        self.assertEqual(
            dict1["state"]["param_1"]["nested_dictionary"]["tensor"].tolist(),
            dict2["state"]["param_1"]["nested_dictionary"]["tensor"].tolist(),
        )
        self.assertEqual(
            dict1["state"]["param_1"]["optimizer_module"]["tensor"].tolist(),
            dict2["state"]["param_1"]["optimizer_module"]["tensor"].tolist(),
        )

        self.assertEqual(
            dict1["state"]["param_1"]["sharded_tensor"].local_shards()[0].tensor.tolist(),
            dict2["state"]["param_1"]["sharded_tensor"].local_shards()[0].tensor.tolist(),
        )
//  [python]
def _assert_state_dict_equals(
        self, dict1: Dict[str, Any], dict2: Dict[str, Any]
    ) -> None:
        self.assertEqual(dict1["param_groups"], dict2["param_groups"])
        self.assertEqual(
            dict1["state"]["param_2"],
            dict2["state"]["param_2"],
        )
        torch.testing.assert_close(
            dict1["state"]["param_1"]["tensor"],
            dict2["state"]["param_1"]["tensor"],
        )
        torch.testing.assert_close(
            dict1["state"]["param_1"]["nested_dictionary"]["tensor"],
            dict2["state"]["param_1"]["nested_dictionary"]["tensor"],
        )
        torch.testing.assert_close(
            dict1["state"]["param_1"]["optimizer_module"]["tensor"],
            dict2["state"]["param_1"]["optimizer_module"]["tensor"],
        )

        torch.testing.assert_close(
            dict1["state"]["param_1"]["sharded_tensor"].local_shards()[0].tensor,
            dict2["state"]["param_1"]["sharded_tensor"].local_shards()[0].tensor,
        )
//  [python]
def check_grads_clipped(self):
        parameters = list(self.parameters())
        assert len(parameters) == len(self.clipped_parameters)
        for actual, expected in zip(parameters, self.clipped_parameters):
            torch.allclose(actual.grad, expected.grad, equal_nan=True)
//  [python]
def check_grads_clipped(self):
        parameters = list(self.parameters())
        assert len(parameters) == len(self.clipped_parameters)
        for actual, expected in zip(parameters, self.clipped_parameters):
            torch.testing.assert_close(actual.grad, expected.grad, equal_nan=True)
//  [python]
def test_scf():
    torch.manual_seed(42)
    nnmodel = (torch.nn.Sequential(torch.nn.Linear(2, 10), torch.nn.Softplus(),
                                   torch.nn.Linear(10, 1, bias=False))).to(
                                       torch.double)
    hybridxc = HybridXC("lda_x", nnmodel, aweight0=0.0)
    e_type = 'ae'
    true_val = '0.09194410469'
    systems = [{
        'moldesc': 'Li 1.5070 0 0; H -1.5070 0 0',
        'basis': '6-311++G(3df,3pd)'
    }, {
        'moldesc': 'Li 0 0 0',
        'basis': '6-311++G(3df,3pd)',
        'spin': 1
    }, {
        'moldesc': 'H 0 0 0',
        'basis': '6-311++G(3df,3pd)',
        'spin': 1
    }]
    entry = DFTEntry.create(e_type, true_val, systems)
    evl = XCNNSCF(hybridxc, entry)
    system = DFTSystem(systems[1])
    run = evl.run(system)
    output = run.energy()
    expected_output = torch.tensor(-7.1914, dtype=torch.float64)
    t.assert_allclose(output, expected_output, atol=1e-4, rtol=0)
//  [python]
def test_scf():
    torch.manual_seed(42)
    nnmodel = (torch.nn.Sequential(torch.nn.Linear(2, 10), torch.nn.Softplus(),
                                   torch.nn.Linear(10, 1, bias=False))).to(
                                       torch.double)
    hybridxc = HybridXC("lda_x", nnmodel, aweight0=0.0)
    e_type = 'ae'
    true_val = '0.09194410469'
    systems = [{
        'moldesc': 'Li 1.5070 0 0; H -1.5070 0 0',
        'basis': '6-311++G(3df,3pd)'
    }, {
        'moldesc': 'Li 0 0 0',
        'basis': '6-311++G(3df,3pd)',
        'spin': 1
    }, {
        'moldesc': 'H 0 0 0',
        'basis': '6-311++G(3df,3pd)',
        'spin': 1
    }]
    entry = DFTEntry.create(e_type, true_val, systems)
    evl = XCNNSCF(hybridxc, entry)
    system = DFTSystem(systems[1])
    run = evl.run(system)
    output = run.energy()
    expected_output = torch.tensor(-7.1914, dtype=torch.float64)
    torch.testing.assert_close(output, expected_output, atol=1e-4, rtol=0)
//  [python]
def check_grads_unscaled(self, optimizer=None):
        if optimizer is not None:
            scaler = self.trainer.precision_plugin.scaler
            state = scaler._per_optimizer_states[id(optimizer)]
            assert state["stage"].name == "UNSCALED"

        grads = [p.grad for p in self.parameters()]
        assert len(grads) == len(self.original_grads)
        for actual, expected in zip(grads, self.original_grads):
            torch.allclose(actual, expected, equal_nan=True)
//  [python]
def check_grads_unscaled(self, optimizer=None):
        if optimizer is not None:
            scaler = self.trainer.precision_plugin.scaler
            state = scaler._per_optimizer_states[id(optimizer)]
            assert state["stage"].name == "UNSCALED"

        grads = [p.grad for p in self.parameters()]
        assert len(grads) == len(self.original_grads)
        for actual, expected in zip(grads, self.original_grads):
            torch.testing.assert_close(actual, expected, equal_nan=True)
//  [python]
def _test_two_configs(
        optim_factory1: Callable[
            [Iterable[Parameter]],
            torch.optim.Optimizer,
        ],
        model_factory1: Callable[
            [torch.device],
            Tuple[
                nn.Module,
                nn.Module,
                torch.Tensor,
                torch.Tensor,
            ],
        ],
        optim_factory2: Callable[
            [Iterable[Parameter]],
            torch.optim.Optimizer,
        ],
        model_factory2: Callable[
            [torch.device],
            Tuple[
                nn.Module,
                nn.Module,
                torch.Tensor,
                torch.Tensor,
            ],
        ],
        device: torch.device,
    ) -> None:
        params1, loss1 = ShampooFSDPDistributorTest._train_model(
            optim_factory1,
            model_factory1,
            device=device,
        )
        params2, loss2 = ShampooFSDPDistributorTest._train_model(
            optim_factory2,
            model_factory2,
            device=device,
        )
        torch.testing.assert_allclose(loss1, loss2)
        torch.testing.assert_allclose(params1, params2)
//  [python]
def _test_two_configs(
        optim_factory1: Callable[
            [Iterable[Parameter]],
            torch.optim.Optimizer,
        ],
        model_factory1: Callable[
            [torch.device],
            Tuple[
                nn.Module,
                nn.Module,
                torch.Tensor,
                torch.Tensor,
            ],
        ],
        optim_factory2: Callable[
            [Iterable[Parameter]],
            torch.optim.Optimizer,
        ],
        model_factory2: Callable[
            [torch.device],
            Tuple[
                nn.Module,
                nn.Module,
                torch.Tensor,
                torch.Tensor,
            ],
        ],
        device: torch.device,
    ) -> None:
        params1, loss1 = ShampooFSDPDistributorTest._train_model(
            optim_factory1,
            model_factory1,
            device=device,
        )
        params2, loss2 = ShampooFSDPDistributorTest._train_model(
            optim_factory2,
            model_factory2,
            device=device,
        )
        torch.testing.assert_close(loss1, loss2)
        torch.testing.assert_close(params1, params2)
//  [python]
def test_scan2d_broadcast(RBLOCK, num_warps, device):

    @triton.jit(debug=True)
    def fn(in_ptr, out_ptr, XBLOCK: tl.constexpr, RBLOCK: tl.constexpr):
        rindex = tl.arange(0, RBLOCK)[None, :]
        xindex = tl.arange(0, XBLOCK)[:, None]
        data = tl.load(in_ptr + rindex)
        scan = tl.cumsum(data, 1)
        expected_max = tl.sum(data, 1)
        tl.device_assert(scan <= expected_max)
        tl.store(out_ptr + xindex * RBLOCK + rindex, scan)

    XBLOCK = 4
    input = torch.randint(0, 10, (1, RBLOCK), dtype=torch.int64, device=device)
    output = torch.empty((XBLOCK, RBLOCK), dtype=torch.int64, device=device)
    fn[(1, )](input, output, XBLOCK, RBLOCK, num_warps=num_warps)
    ref = input.cumsum(1).broadcast_to((XBLOCK, RBLOCK))
    torch.testing.assert_allclose(output, ref)
//  [python]
def test_scan2d_broadcast(RBLOCK, num_warps, device):

    @triton.jit(debug=True)
    def fn(in_ptr, out_ptr, XBLOCK: tl.constexpr, RBLOCK: tl.constexpr):
        rindex = tl.arange(0, RBLOCK)[None, :]
        xindex = tl.arange(0, XBLOCK)[:, None]
        data = tl.load(in_ptr + rindex)
        scan = tl.cumsum(data, 1)
        expected_max = tl.sum(data, 1)
        tl.device_assert(scan <= expected_max)
        tl.store(out_ptr + xindex * RBLOCK + rindex, scan)

    XBLOCK = 4
    input = torch.randint(0, 10, (1, RBLOCK), dtype=torch.int64, device=device)
    output = torch.empty((XBLOCK, RBLOCK), dtype=torch.int64, device=device)
    fn[(1, )](input, output, XBLOCK, RBLOCK, num_warps=num_warps)
    ref = input.cumsum(1).broadcast_to((XBLOCK, RBLOCK))
    torch.testing.assert_close(output, ref)
//  [python]
def test_scan2d_for(device):

    @triton.jit
    def fn(out_ptr0, rnumel, RBLOCK: tl.constexpr):
        rbase = tl.arange(0, RBLOCK)[None, :]
        for roffset in range(0, rnumel, RBLOCK):
            rindex = roffset + rbase
            rmask = rindex < rnumel
            tmp3 = tl.where(rmask, 1, 0)
            tmp6 = tl.cumsum(tmp3, 1)
            tl.store(out_ptr0 + rindex, tmp6, rmask)

    RBLOCK = 8
    out0 = torch.empty(RBLOCK, device=device, dtype=torch.int64)
    fn[(1, )](out0, RBLOCK, RBLOCK)
    ref = torch.arange(RBLOCK, device=device, dtype=torch.int64) + 1
    torch.testing.assert_allclose(out0, ref)
//  [python]
def test_scan2d_for(device):

    @triton.jit
    def fn(out_ptr0, rnumel, RBLOCK: tl.constexpr):
        rbase = tl.arange(0, RBLOCK)[None, :]
        for roffset in range(0, rnumel, RBLOCK):
            rindex = roffset + rbase
            rmask = rindex < rnumel
            tmp3 = tl.where(rmask, 1, 0)
            tmp6 = tl.cumsum(tmp3, 1)
            tl.store(out_ptr0 + rindex, tmp6, rmask)

    RBLOCK = 8
    out0 = torch.empty(RBLOCK, device=device, dtype=torch.int64)
    fn[(1, )](out0, RBLOCK, RBLOCK)
    ref = torch.arange(RBLOCK, device=device, dtype=torch.int64) + 1
    torch.testing.assert_close(out0, ref)
//  [python]
def test_as_torch_tensor():
    """
    test __torch_func__ protocol
    """
    a = hidet.randn([32, 32], dtype='float16', device='cuda')
    b = torch.abs(a)
    c = hidet.ops.abs(a)
    torch.testing.assert_allclose(b, c.torch())
//  [python]
def test_as_torch_tensor():
    """
    test __torch_func__ protocol
    """
    a = hidet.randn([32, 32], dtype='float16', device='cuda')
    b = torch.abs(a)
    c = hidet.ops.abs(a)
    torch.testing.assert_close(b, c.torch())
//  [python]
def _test(targets: List[float], scheduler: LRScheduler, epochs: int, optimizer: Optimizer, ssr: float):
        scale_pytorch_scheduler(scheduler, ssr)
        for epoch in range(epochs):
            for param_group in optimizer.param_groups:
                t.assert_close(targets[epoch], param_group['lr'])
            scheduler.step()
//  [python]
def _test(targets: List[float], scheduler: LRScheduler, epochs: int, optimizer: Optimizer, ssr: float):
        scale_pytorch_scheduler(scheduler, ssr)
        for epoch in range(epochs):
            for param_group in optimizer.param_groups:
                torch.testing.assert_close(targets[epoch], param_group['lr'])
            scheduler.step()
//  [python]
def test_evaluate(tmp_path, trainer_kwargs):
    dm = ClassifDataModule()
    model = ClassificationModel()
    trainer = Trainer(
        default_root_dir=tmp_path, max_epochs=2, limit_train_batches=10, limit_val_batches=10, **trainer_kwargs
    )

    trainer.fit(model, datamodule=dm)
    assert "ckpt" in trainer.checkpoint_callback.best_model_path

    old_weights = model.layer_0.weight.clone().detach().cpu()

    trainer.validate(datamodule=dm)
    trainer.test(datamodule=dm)

    # make sure weights didn't change
    new_weights = model.layer_0.weight.clone().detach().cpu()
    t.assert_allclose(old_weights, new_weights)
//  [python]
def test_evaluate(tmp_path, trainer_kwargs):
    dm = ClassifDataModule()
    model = ClassificationModel()
    trainer = Trainer(
        default_root_dir=tmp_path, max_epochs=2, limit_train_batches=10, limit_val_batches=10, **trainer_kwargs
    )

    trainer.fit(model, datamodule=dm)
    assert "ckpt" in trainer.checkpoint_callback.best_model_path

    old_weights = model.layer_0.weight.clone().detach().cpu()

    trainer.validate(datamodule=dm)
    trainer.test(datamodule=dm)

    # make sure weights didn't change
    new_weights = model.layer_0.weight.clone().detach().cpu()
    torch.testing.assert_close(old_weights, new_weights)
//  [python]
def assert_tensor_equal_over_group(tensor: torch.Tensor, group: ProcessGroup, assert_: bool = True) -> bool:
    """We assume that tensors are already of correct size."""
    reference_rank = 0
    if dist.get_rank(group) == reference_rank:
        reference_tensor = tensor
    else:
        reference_tensor = torch.empty_like(tensor)
    dist.broadcast(
        reference_tensor,
        src=get_global_rank(group=group, group_rank=reference_rank),
        group=group,
    )
    if assert_:
        torch.allclose(tensor, reference_tensor, atol=0, rtol=0)
    else:
        result = torch.allclose(tensor, reference_tensor, atol=0.0, rtol=0.0)
        results = [0] * group.size()
        dist.all_gather_object(results, result, group)
        return all(results)
//  [python]
def assert_tensor_equal_over_group(tensor: torch.Tensor, group: ProcessGroup, assert_: bool = True) -> bool:
    """We assume that tensors are already of correct size."""
    reference_rank = 0
    if dist.get_rank(group) == reference_rank:
        reference_tensor = tensor
    else:
        reference_tensor = torch.empty_like(tensor)
    dist.broadcast(
        reference_tensor,
        src=get_global_rank(group=group, group_rank=reference_rank),
        group=group,
    )
    if assert_:
        torch.testing.assert_close(tensor, reference_tensor, atol=0, rtol=0)
    else:
        result = torch.allclose(tensor, reference_tensor, atol=0.0, rtol=0.0)
        results = [0] * group.size()
        dist.all_gather_object(results, result, group)
        return all(results)
//  [python]
def test_inv_remap_state_dict(model_name: str):
    """
    Verify that we can convert a HF BigCode model to flash_attn and back.
    """

    state_dict = state_dict_from_pretrained(model_name)
    config = GPTBigCodeConfig.from_pretrained(model_name)

    flash_state_dict = remap_state_dict_hf_bigcode(state_dict, config)
    recovered_state_dict = inv_remap_state_dict_hf_bigcode(flash_state_dict, config)

    assert set(state_dict.keys()) == set(recovered_state_dict.keys())

    for k in state_dict.keys():
        assert state_dict[k].shape == recovered_state_dict[k].shape
        t.assert_allclose(state_dict[k], recovered_state_dict[k], rtol=1e-6, atol=1e-6)
//  [python]
def test_inv_remap_state_dict(model_name: str):
    """
    Verify that we can convert a HF BigCode model to flash_attn and back.
    """

    state_dict = state_dict_from_pretrained(model_name)
    config = GPTBigCodeConfig.from_pretrained(model_name)

    flash_state_dict = remap_state_dict_hf_bigcode(state_dict, config)
    recovered_state_dict = inv_remap_state_dict_hf_bigcode(flash_state_dict, config)

    assert set(state_dict.keys()) == set(recovered_state_dict.keys())

    for k in state_dict.keys():
        assert state_dict[k].shape == recovered_state_dict[k].shape
        torch.testing.assert_close(state_dict[k], recovered_state_dict[k], rtol=1e-6, atol=1e-6)
//  [python]
def _test_cast(
        self,
        shape,
        dtype="float32",
        cast_dtype="bfloat16",
        test_name="cast",
    ) -> None:
        if not isinstance(shape, list):
            shape = [shape]

        X = Tensor(
            shape=shape,
            name="X",
            dtype=dtype,
            is_input=True,
        )

        Y = ops.cast()(X, cast_dtype)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        module = compile_model(Y, target, "./tmp", f"{test_name}_{self._test_id}")
        self._test_id += 1

        x = get_random_torch_tensor(shape, dtype=dtype)
        y = get_torch_empty_tensor(shape, dtype=cast_dtype)
        inputs = {"X": x}
        outputs = {"Y": y}
        module.run_with_tensors(inputs, outputs)

        y_pt = x.to(string_to_torch_dtype(cast_dtype))
        t.assert_close(y, y_pt, atol=1e-2, rtol=1e-2)
//  [python]
def _test_cast(
        self,
        shape,
        dtype="float32",
        cast_dtype="bfloat16",
        test_name="cast",
    ) -> None:
        if not isinstance(shape, list):
            shape = [shape]

        X = Tensor(
            shape=shape,
            name="X",
            dtype=dtype,
            is_input=True,
        )

        Y = ops.cast()(X, cast_dtype)
        Y._attrs["name"] = "Y"
        Y._attrs["is_output"] = True

        target = detect_target()
        module = compile_model(Y, target, "./tmp", f"{test_name}_{self._test_id}")
        self._test_id += 1

        x = get_random_torch_tensor(shape, dtype=dtype)
        y = get_torch_empty_tensor(shape, dtype=cast_dtype)
        inputs = {"X": x}
        outputs = {"Y": y}
        module.run_with_tensors(inputs, outputs)

        y_pt = x.to(string_to_torch_dtype(cast_dtype))
        torch.testing.assert_close(y, y_pt, atol=1e-2, rtol=1e-2)
//  [python]
def test_elementwise_model(self):
        def run_elementwise_model():
            # A function to test.
            def elementwise_model(x):
                w = x.relu()
                y = w * w + 1.5
                z = y + x
                p = z * x
                q = p.relu()
                return q

            def run(fun, device, x):
                x = torch.tensor(x, device=device, dtype=torch.float32).requires_grad_()
                y = fun(x)
                y.sum().backward()
                return x, y, x.grad

            # Baseline.
            x, y, g_x = run(elementwise_model, "cpu", [-1.0, 2.0])
            # ORT result.
            x_new, y_new, g_x_new = run(elementwise_model, "lazy", [-1.0, 2.0])

            torch.testing.assert_allclose(x.to("lazy"), x_new)
            torch.testing.assert_allclose(y.to("lazy"), y_new)
            torch.testing.assert_allclose(g_x.to("lazy"), g_x_new)

        for _ in range(5):
            run_elementwise_model()
//  [python]
def test_elementwise_model(self):
        def run_elementwise_model():
            # A function to test.
            def elementwise_model(x):
                w = x.relu()
                y = w * w + 1.5
                z = y + x
                p = z * x
                q = p.relu()
                return q

            def run(fun, device, x):
                x = torch.tensor(x, device=device, dtype=torch.float32).requires_grad_()
                y = fun(x)
                y.sum().backward()
                return x, y, x.grad

            # Baseline.
            x, y, g_x = run(elementwise_model, "cpu", [-1.0, 2.0])
            # ORT result.
            x_new, y_new, g_x_new = run(elementwise_model, "lazy", [-1.0, 2.0])

            torch.testing.assert_close(x.to("lazy"), x_new)
            torch.testing.assert_close(y.to("lazy"), y_new)
            torch.testing.assert_close(g_x.to("lazy"), g_x_new)

        for _ in range(5):
            run_elementwise_model()
//  [python]
def test_load_state_dict(self) -> None:
        def get_optimizer() -> WarmupOptimizer:
            param_1_t = torch.tensor([1.0, 2.0])
            param_1 = Variable(param_1_t)
            keyed_optimizer = DummyKeyedOptimizer(
                {"param_1": param_1}, defaultdict(dict), [{"params": [param_1]}]
            )
            warmup_optimizer = WarmupOptimizer(
                keyed_optimizer,
                stages=[
                    WarmupStage(
                        WarmupPolicy.LINEAR, max_iters=100, value=1e-2, lr_scale=1
                    ),
                ],
            )
            warmup_optimizer.save_param_groups(True)
            return warmup_optimizer

        warmup_optimizer_1 = get_optimizer()
        num_iters = 10
        for _ in range(num_iters):
            warmup_optimizer_1.zero_grad()
            warmup_optimizer_1.step()

        param_state = list(warmup_optimizer_1.state.values())[0]
        self.assertEqual(
            param_state["warmup"].tolist()[0],
            num_iters,
        )

        warmup_optimizer_2 = get_optimizer()
        warmup_optimizer_2.step()
        warmup_optimizer_2.zero_grad()

        warmup_optimizer_2.save_param_groups(True)
        warmup_optimizer_2.load_state_dict(warmup_optimizer_1.state_dict())

        self.assertEqual(
            warmup_optimizer_1.state_dict()["param_groups"],
            warmup_optimizer_2.state_dict()["param_groups"],
        )
        self.assertTrue(
            torch.allclose(
                warmup_optimizer_1.state_dict()["state"]["__warmup"],
                warmup_optimizer_2.state_dict()["state"]["__warmup"],
            )
        )
//  [python]
def test_load_state_dict(self) -> None:
        def get_optimizer() -> WarmupOptimizer:
            param_1_t = torch.tensor([1.0, 2.0])
            param_1 = Variable(param_1_t)
            keyed_optimizer = DummyKeyedOptimizer(
                {"param_1": param_1}, defaultdict(dict), [{"params": [param_1]}]
            )
            warmup_optimizer = WarmupOptimizer(
                keyed_optimizer,
                stages=[
                    WarmupStage(
                        WarmupPolicy.LINEAR, max_iters=100, value=1e-2, lr_scale=1
                    ),
                ],
            )
            warmup_optimizer.save_param_groups(True)
            return warmup_optimizer

        warmup_optimizer_1 = get_optimizer()
        num_iters = 10
        for _ in range(num_iters):
            warmup_optimizer_1.zero_grad()
            warmup_optimizer_1.step()

        param_state = list(warmup_optimizer_1.state.values())[0]
        self.assertEqual(
            param_state["warmup"].tolist()[0],
            num_iters,
        )

        warmup_optimizer_2 = get_optimizer()
        warmup_optimizer_2.step()
        warmup_optimizer_2.zero_grad()

        warmup_optimizer_2.save_param_groups(True)
        warmup_optimizer_2.load_state_dict(warmup_optimizer_1.state_dict())

        self.assertEqual(
            warmup_optimizer_1.state_dict()["param_groups"],
            warmup_optimizer_2.state_dict()["param_groups"],
        )
        torch.testing.assert_close(
            warmup_optimizer_1.state_dict()["state"]["__warmup"],
            warmup_optimizer_2.state_dict()["state"]["__warmup"],
        )
//  [python]
def test_inference_logits(self):
        r"""
        Logits testing to check implementation consistency between `fairseq` implementation
        and `transformers` implementation of NLLB-MoE transformers. We only check the logits
        of the second sample of the batch, as it is padded.
        """
        model = NllbMoeForConditionalGeneration.from_pretrained("hf-internal-testing/random-nllb-moe-2-experts").eval()
        with torch.no_grad():
            output = model(**self.model_inputs)

        EXPECTED_LOGTIS = torch.Tensor([-0.3059, 0.0000, 9.3029, 0.6456, -0.9148, 1.7836, 0.6478, 0.9438, -0.5272, -0.6617, -1.2717, 0.4564, 0.1345, -0.2301, -1.0140, 1.1427, -1.5535, 0.1337, 0.2082, -0.8112, -0.3842, -0.3377, 0.1256, 0.6450, -0.0452, 0.0219, 1.4274, -0.4991, -0.2063, -0.4409,])  # fmt: skip
        torch.allclose(output.logits[1, 0, :30], EXPECTED_LOGTIS, rtol=6e-3, atol=9e-3)
//  [python]
def test_cosine_sim_torch():
    a = torch.rand(128)
    b = torch.rand(128)
    assert metrics.cosine_sim(a, b).shape == (1,)
    assert metrics.cosine_sim(a, b) == metrics.cosine_sim(b, a)
    torch.testing.assert_allclose(metrics.cosine_sim(a, a), torch.ones(1))

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    assert metrics.cosine_sim(a, b).shape == (10, 5)
    assert metrics.cosine_sim(b, a).shape == (5, 10)
    diag_dists = torch.diagonal(metrics.cosine_sim(b, b))  # self-comparisons
    torch.testing.assert_allclose(diag_dists, torch.ones(5))
//  [python]
def test_cosine_sim_torch():
    a = torch.rand(128)
    b = torch.rand(128)
    assert metrics.cosine_sim(a, b).shape == (1,)
    assert metrics.cosine_sim(a, b) == metrics.cosine_sim(b, a)
    torch.testing.assert_close(metrics.cosine_sim(a, a), torch.ones(1))

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    assert metrics.cosine_sim(a, b).shape == (10, 5)
    assert metrics.cosine_sim(b, a).shape == (5, 10)
    diag_dists = torch.diagonal(metrics.cosine_sim(b, b))  # self-comparisons
    torch.testing.assert_allclose(diag_dists, torch.ones(5))
//  [python]
def test_sqeuclidean_dist_torch():
    a = torch.rand(128)
    b = torch.rand(128)
    assert metrics.sqeuclidean_dist(a, b).shape == (1,)
    torch.testing.assert_allclose(
        metrics.sqeuclidean_dist(a, b),
        metrics.euclidean_dist(a, b) ** 2,
    )

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    assert metrics.sqeuclidean_dist(a, b).shape == (10, 5)
    torch.testing.assert_allclose(
        metrics.sqeuclidean_dist(a, b),
        metrics.euclidean_dist(a, b) ** 2,
    )
//  [python]
def test_sqeuclidean_dist_torch():
    a = torch.rand(128)
    b = torch.rand(128)
    assert metrics.sqeuclidean_dist(a, b).shape == (1,)
    torch.testing.assert_close(
        metrics.sqeuclidean_dist(a, b),
        metrics.euclidean_dist(a, b) ** 2,
    )

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    assert metrics.sqeuclidean_dist(a, b).shape == (10, 5)
    torch.testing.assert_close(
        metrics.sqeuclidean_dist(a, b),
        metrics.euclidean_dist(a, b) ** 2,
    )
//  [python]
def group_gemm_test_helper(self, compute_type, count:int, ms:List[int], ns:List[int], ks:List[int], rtol:float,
                               atol:float):

        As_ref = [random_tensor((m, k), dtype=compute_type, device="cuda") for m, k in zip(ms, ks)]
        Bs_ref = [random_tensor((k, n), dtype=compute_type, device="cuda") for k, n in zip(ks, ns)]
        Cs_ref = [torch.zeros((m, n), dtype=compute_type, device="cuda") for m, n in zip(ms, ns)]
        
        self.starter.record()
        FT_result = self.group_gemm(As_ref, Bs_ref)
        self.ender.record()
        torch.cuda.synchronize()
        curr_time = self.starter.elapsed_time(self.ender)
        print(f"group gemm time is {curr_time}")

        for i in range(count):
            Cs_ref[i] = torch.mm(As_ref[i], Bs_ref[i])
            self.assertEqual(FT_result[i], Cs_ref[i], rtol=rtol, atol=atol)
//  [python]
def group_gemm_test_helper(self, compute_type, count:int, ms:List[int], ns:List[int], ks:List[int], rtol:float,
                               atol:float):

        As_ref = [random_tensor((m, k), dtype=compute_type, device="cuda") for m, k in zip(ms, ks)]
        Bs_ref = [random_tensor((k, n), dtype=compute_type, device="cuda") for k, n in zip(ks, ns)]
        Cs_ref = [torch.zeros((m, n), dtype=compute_type, device="cuda") for m, n in zip(ms, ns)]
        
        self.starter.record()
        FT_result = self.group_gemm(As_ref, Bs_ref)
        self.ender.record()
        torch.cuda.synchronize()
        curr_time = self.starter.elapsed_time(self.ender)
        print(f"group gemm time is {curr_time}")

        for i in range(count):
            Cs_ref[i] = torch.mm(As_ref[i], Bs_ref[i])
            torch.testing.assert_close(FT_result[i], Cs_ref[i], rtol=rtol, atol=atol, check_dtype=False)
//  [python]
def test_euclidean_dist_torch():
    a = torch.rand(128)
    b = torch.rand(128)
    assert metrics.euclidean_dist(a, b).shape == (1,)
    assert metrics.euclidean_dist(a, b) == metrics.euclidean_dist(b, a)
    torch.testing.assert_allclose(metrics.euclidean_dist(a, a), torch.zeros(1))

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    assert metrics.euclidean_dist(a, b).shape == (10, 5)
    assert metrics.euclidean_dist(b, a).shape == (5, 10)
    diag_dists = torch.diagonal(metrics.euclidean_dist(b, b))  # self-comparisons
    torch.testing.assert_allclose(diag_dists, torch.zeros(5))

    a = torch.tensor([0.0, 2.0, 0.0])
    b = torch.tensor([0.0, 0.0, 2.0])
    desired_output_singleton = torch.sqrt(torch.tensor([2.0**2.0 + 2.0**2.0]))
    torch.testing.assert_allclose(metrics.euclidean_dist(a, b), desired_output_singleton)

    a = torch.tensor([[0.0, 2.0, 0.0], [0.0, 0.0, 2.0]])
    b = torch.tensor([[0.0, 0.0, 2.0], [0.0, 2.0, 0.0]])
    desired_output_singleton = torch.tensor(
        [[desired_output_singleton.item(), 0.0], [0.0, desired_output_singleton.item()]]
    )
    torch.testing.assert_allclose(metrics.euclidean_dist(a, b), desired_output_singleton)
//  [python]
def test_euclidean_dist_torch():
    a = torch.rand(128)
    b = torch.rand(128)
    assert metrics.euclidean_dist(a, b).shape == (1,)
    assert metrics.euclidean_dist(a, b) == metrics.euclidean_dist(b, a)
    torch.testing.assert_close(metrics.euclidean_dist(a, a), torch.zeros(1))

    a = torch.rand(10, 3)
    b = torch.rand(5, 3)
    assert metrics.euclidean_dist(a, b).shape == (10, 5)
    assert metrics.euclidean_dist(b, a).shape == (5, 10)
    diag_dists = torch.diagonal(metrics.euclidean_dist(b, b))  # self-comparisons
    torch.testing.assert_allclose(diag_dists, torch.zeros(5))

    a = torch.tensor([0.0, 2.0, 0.0])
    b = torch.tensor([0.0, 0.0, 2.0])
    desired_output_singleton = torch.sqrt(torch.tensor([2.0**2.0 + 2.0**2.0]))
    torch.testing.assert_close(metrics.euclidean_dist(a, b), desired_output_singleton)

    a = torch.tensor([[0.0, 2.0, 0.0], [0.0, 0.0, 2.0]])
    b = torch.tensor([[0.0, 0.0, 2.0], [0.0, 2.0, 0.0]])
    desired_output_singleton = torch.tensor(
        [[desired_output_singleton.item(), 0.0], [0.0, desired_output_singleton.item()]]
    )
    torch.testing.assert_close(metrics.euclidean_dist(a, b), desired_output_singleton)
//  [python]
def test_forward_backward(
    shape,
    dtype,
):
    m_tiles, n_tiles, k_tiles = shape
    m, n, k = sum(m_tiles), sum(n_tiles), sum(k_tiles)

    torch.manual_seed(m * n * k)

    a, b, c_reference = make_operands(m, n, k, dtype=dtype)
    a = a.cuda().requires_grad_()
    b = b.cuda().requires_grad_()
    c_reference = c_reference.cuda()

    # In one operand make each tile have its own strides, in the other use the
    # same stride for all tiles. And make the two operands have the stride==1
    # in different dimensions.
    a_tiled = [
        [y.t().clone().t() for y in x.split(k_tiles, dim=1)]
        for x in a.split(m_tiles, dim=0)
    ]
    b_tiled = [[y for y in x.split(n_tiles, dim=1)] for x in b.split(k_tiles, dim=0)]

    c_test_tiled = tiled_matmul(a_tiled, b_tiled)
    c_test = torch.cat([torch.cat(x, dim=1) for x in c_test_tiled], dim=0)

    torch.testing.assert_allclose(c_test, c_reference)

    # To avoid numerical issues in the backward, set things up so that we only
    # multiply by a diagonal matrix whose entries are +/- 2^{-1/0/+1} (so that
    # it only changes the sign bit and the exponent).
    diag = torch.tensor(random.choices([-2, -1, -0.5, 0.5, 1, 2], k=min(m, n)))
    grad_c = torch.zeros_like(c_test)
    torch.diag(grad_c)[:] = diag
    grad_a_reference = torch.matmul(grad_c, b.detach().t())
    grad_b_reference = torch.matmul(a.detach().t(), grad_c)

    torch.autograd.backward([c_test], [grad_c], inputs=[a, b])

    torch.testing.assert_allclose(a.grad, grad_a_reference)
    torch.testing.assert_allclose(b.grad, grad_b_reference)
//  [python]
def test_forward_backward(
    shape,
    dtype,
):
    m_tiles, n_tiles, k_tiles = shape
    m, n, k = sum(m_tiles), sum(n_tiles), sum(k_tiles)

    torch.manual_seed(m * n * k)

    a, b, c_reference = make_operands(m, n, k, dtype=dtype)
    a = a.cuda().requires_grad_()
    b = b.cuda().requires_grad_()
    c_reference = c_reference.cuda()

    # In one operand make each tile have its own strides, in the other use the
    # same stride for all tiles. And make the two operands have the stride==1
    # in different dimensions.
    a_tiled = [
        [y.t().clone().t() for y in x.split(k_tiles, dim=1)]
        for x in a.split(m_tiles, dim=0)
    ]
    b_tiled = [[y for y in x.split(n_tiles, dim=1)] for x in b.split(k_tiles, dim=0)]

    c_test_tiled = tiled_matmul(a_tiled, b_tiled)
    c_test = torch.cat([torch.cat(x, dim=1) for x in c_test_tiled], dim=0)

    torch.testing.assert_close(c_test, c_reference)

    # To avoid numerical issues in the backward, set things up so that we only
    # multiply by a diagonal matrix whose entries are +/- 2^{-1/0/+1} (so that
    # it only changes the sign bit and the exponent).
    diag = torch.tensor(random.choices([-2, -1, -0.5, 0.5, 1, 2], k=min(m, n)))
    grad_c = torch.zeros_like(c_test)
    torch.diag(grad_c)[:] = diag
    grad_a_reference = torch.matmul(grad_c, b.detach().t())
    grad_b_reference = torch.matmul(a.detach().t(), grad_c)

    torch.autograd.backward([c_test], [grad_c], inputs=[a, b])

    torch.testing.assert_close(a.grad, grad_a_reference)
    torch.testing.assert_close(b.grad, grad_b_reference)
//  [python]
def assert_all_approx_close(a, b, rtol=1e-3, atol=1e-3, count=0, throw=True):
    idx = torch.isclose(a, b, rtol=rtol, atol=atol)
    sumval = (idx == 0).sum().item()
    if sumval > count:
        if throw:
            print(f"Too many values not close: assert {sumval} < {count}")
            t.assert_close(a, b, rtol=rtol, atol=atol)

    return sumval
//  [python]
def assert_all_approx_close(a, b, rtol=1e-3, atol=1e-3, count=0, throw=True):
    idx = torch.isclose(a, b, rtol=rtol, atol=atol)
    sumval = (idx == 0).sum().item()
    if sumval > count:
        if throw:
            print(f"Too many values not close: assert {sumval} < {count}")
            torch.testing.assert_close(a, b, rtol=rtol, atol=atol)

    return sumval
//  [python]
def test_estimate_quantiles(dtype):
    A = torch.rand(1024, 1024, device="cuda")
    A = A.to(dtype)
    code = F.estimate_quantiles(A)

    percs = torch.linspace(1 / 512, 511 / 512, 256, device=A.device)
    torch.testing.assert_allclose(percs, code, atol=1e-3, rtol=1e-2)

    A = torch.randn(1024, 1024, device="cuda")
    A = A.to(dtype)
    code = F.estimate_quantiles(A)

    quantiles = torch.quantile(A.float(), percs)
    diff = torch.abs(code - quantiles)
    assert (diff > 5e-02).sum().item() == 0
//  [python]
def test_estimate_quantiles(dtype):
    A = torch.rand(1024, 1024, device="cuda")
    A = A.to(dtype)
    code = F.estimate_quantiles(A)

    percs = torch.linspace(1 / 512, 511 / 512, 256, device=A.device)
    torch.testing.assert_close(percs, code, atol=1e-3, rtol=1e-2)

    A = torch.randn(1024, 1024, device="cuda")
    A = A.to(dtype)
    code = F.estimate_quantiles(A)

    quantiles = torch.quantile(A.float(), percs)
    diff = torch.abs(code - quantiles)
    assert (diff > 5e-02).sum().item() == 0
//  [python]
def test_get_model_state_dict(self):
        device = torch.device("cuda", self.rank)

        with device:
            model = MyModel(4)
        expected = model.state_dict()
        assert not has_fsdp_modules(model)

        # No sharding - full state dict
        options = StateDictOptions(full_state_dict=True, cpu_offload=False)
        state_dict = get_model_state_dict(model, options, self.rank)
        torch.testing.assert_allclose(state_dict, expected)

        options = StateDictOptions(full_state_dict=True, cpu_offload=True)
        state_dict = get_model_state_dict(model, options, self.rank)
        expected_cpu = {k: v.to(device="cpu") for k, v in expected.items()}
        torch.testing.assert_allclose(state_dict, expected_cpu)

        options = StateDictOptions(full_state_dict=True, rank0_only=True)
        state_dict = get_model_state_dict(model, options, self.rank)
        torch.testing.assert_allclose(state_dict, expected if self.rank == 0 else {})

        # No sharding - sharded state dict
        options = StateDictOptions(full_state_dict=False)
        with pytest.raises(ValueError, match="cannot be used"):
            get_model_state_dict(model, options, self.rank)

        sharded_model = thunder.distributed.fsdp(model)
        assert has_fsdp_modules(model)

        # Sharding - full state dict
        options = StateDictOptions(full_state_dict=True, cpu_offload=False)
        state_dict = get_model_state_dict(sharded_model, options, self.rank)
        torch.testing.assert_allclose(state_dict, expected)

        options = StateDictOptions(full_state_dict=True, cpu_offload=True)
        state_dict = get_model_state_dict(sharded_model, options, self.rank)
        torch.testing.assert_allclose(state_dict, expected_cpu)

        options = StateDictOptions(full_state_dict=True, rank0_only=True)
        state_dict = get_model_state_dict(sharded_model, options, self.rank)
        torch.testing.assert_allclose(state_dict, expected if self.rank == 0 else {})

        tmp_path = self.tmp_path

        # Sharding - sharded state dict
        for value in (False, True):
            options = StateDictOptions(full_state_dict=False, cpu_offload=value)
            state_dict = get_model_state_dict(sharded_model, options, self.rank)

            assert isinstance(state_dict["buf"], torch.Tensor)
            assert isinstance(state_dict["l.weight"], DTensor)
            assert state_dict["buf"].device.type == ("cpu" if value else "cuda")
            assert state_dict["l.weight"].to_local().device.type == "cuda"

            save(state_dict, tmp_path)
            torch.distributed.barrier()
            assert set(os.listdir(tmp_path)) == {"__1_0.distcp", "__0_0.distcp", ".metadata"}
            checkpoint = distributed_ckpt_to_regular(tmp_path)
            if self.rank == 0:
                torch.testing.assert_allclose(checkpoint, expected_cpu)

        options = StateDictOptions(full_state_dict=False, rank0_only=True)
        with pytest.raises(ValueError, match="cannot be used"):
            get_model_state_dict(sharded_model, options, self.rank)
//  [python]
def test_get_model_state_dict(self):
        device = torch.device("cuda", self.rank)

        with device:
            model = MyModel(4)
        expected = model.state_dict()
        assert not has_fsdp_modules(model)

        # No sharding - full state dict
        options = StateDictOptions(full_state_dict=True, cpu_offload=False)
        state_dict = get_model_state_dict(model, options, self.rank)
        torch.testing.assert_close(state_dict, expected)

        options = StateDictOptions(full_state_dict=True, cpu_offload=True)
        state_dict = get_model_state_dict(model, options, self.rank)
        expected_cpu = {k: v.to(device="cpu") for k, v in expected.items()}
        torch.testing.assert_close(state_dict, expected_cpu)

        options = StateDictOptions(full_state_dict=True, rank0_only=True)
        state_dict = get_model_state_dict(model, options, self.rank)
        torch.testing.assert_close(state_dict, expected if self.rank == 0 else {})

        # No sharding - sharded state dict
        options = StateDictOptions(full_state_dict=False)
        with pytest.raises(ValueError, match="cannot be used"):
            get_model_state_dict(model, options, self.rank)

        sharded_model = thunder.distributed.fsdp(model)
        assert has_fsdp_modules(model)

        # Sharding - full state dict
        options = StateDictOptions(full_state_dict=True, cpu_offload=False)
        state_dict = get_model_state_dict(sharded_model, options, self.rank)
        torch.testing.assert_close(state_dict, expected)

        options = StateDictOptions(full_state_dict=True, cpu_offload=True)
        state_dict = get_model_state_dict(sharded_model, options, self.rank)
        torch.testing.assert_close(state_dict, expected_cpu)

        options = StateDictOptions(full_state_dict=True, rank0_only=True)
        state_dict = get_model_state_dict(sharded_model, options, self.rank)
        torch.testing.assert_close(state_dict, expected if self.rank == 0 else {})

        tmp_path = self.tmp_path

        # Sharding - sharded state dict
        for value in (False, True):
            options = StateDictOptions(full_state_dict=False, cpu_offload=value)
            state_dict = get_model_state_dict(sharded_model, options, self.rank)

            assert isinstance(state_dict["buf"], torch.Tensor)
            assert isinstance(state_dict["l.weight"], DTensor)
            assert state_dict["buf"].device.type == ("cpu" if value else "cuda")
            assert state_dict["l.weight"].to_local().device.type == "cuda"

            save(state_dict, tmp_path)
            torch.distributed.barrier()
            assert set(os.listdir(tmp_path)) == {"__1_0.distcp", "__0_0.distcp", ".metadata"}
            checkpoint = distributed_ckpt_to_regular(tmp_path)
            if self.rank == 0:
                torch.testing.assert_close(checkpoint, expected_cpu)

        options = StateDictOptions(full_state_dict=False, rank0_only=True)
        with pytest.raises(ValueError, match="cannot be used"):
            get_model_state_dict(sharded_model, options, self.rank)
//  [python]
def test_load_model_state_dict(self):
        device = torch.device("cuda", self.rank)
        pg = c10d.new_group()

        state_dict = MyModel(4).to(device=device).state_dict()

        # No sharding - full state dict
        model = MyModel(4).to(device=device)
        options = StateDictOptions(full_state_dict=True, cpu_offload=False)
        load_model_state_dict(state_dict, model, options, self.rank)
        t.assert_close(model.state_dict(), state_dict)

        # cpu_offload=True is not relevant for this case

        model = MyModel(4).to(device=device)
        options = StateDictOptions(full_state_dict=True, rank0_only=True)
        load_model_state_dict(state_dict, model, options, self.rank)
        if self.rank == 0:
            t.assert_close(model.state_dict(), state_dict)
        else:
            with pytest.raises(AssertionError):
                t.assert_close(model.state_dict(), state_dict)

        # No sharding - sharded state dict
        options = StateDictOptions(full_state_dict=False)
        with pytest.raises(ValueError, match="cannot be used"):
            load_model_state_dict(state_dict, model, options, self.rank)

        # Sharding - full state dict
        for kwargs in ({"cpu_offload": True}, {"cpu_offload": False}, {"rank0_only": True}):
            model = MyModel(4).to(device=device)
            sharded_model = thunder.distributed.fsdp(model)
            options = StateDictOptions(full_state_dict=True, **kwargs)
            load_model_state_dict(state_dict, sharded_model, options, self.rank)
            _unshard_params(sharded_model, pg)
            t.assert_close(model.state_dict(), state_dict)

        # Create a sharded state_dict that can be loaded
        model = MyModel(4).to(device=device)
        sharded_model_expected = thunder.distributed.fsdp(model)
        options = StateDictOptions(full_state_dict=False)
        sharded_state_dict = get_model_state_dict(sharded_model_expected, options, self.rank)

        # Sharding - sharded state dict
        for kwargs in ({"cpu_offload": True}, {"cpu_offload": False}, {"rank0_only": True}):
            model = MyModel(4).to(device=device)
            sharded_model = thunder.distributed.fsdp(model)
            options = StateDictOptions(full_state_dict=False, **kwargs)
            load_model_state_dict(sharded_state_dict, sharded_model, options, self.rank)
            t.assert_close(sharded_model.state_dict(), sharded_model_expected.state_dict())
//  [python]
def test_load_model_state_dict(self):
        device = torch.device("cuda", self.rank)
        pg = c10d.new_group()

        state_dict = MyModel(4).to(device=device).state_dict()

        # No sharding - full state dict
        model = MyModel(4).to(device=device)
        options = StateDictOptions(full_state_dict=True, cpu_offload=False)
        load_model_state_dict(state_dict, model, options, self.rank)
        torch.testing.assert_close(model.state_dict(), state_dict)

        # cpu_offload=True is not relevant for this case

        model = MyModel(4).to(device=device)
        options = StateDictOptions(full_state_dict=True, rank0_only=True)
        load_model_state_dict(state_dict, model, options, self.rank)
        if self.rank == 0:
            torch.testing.assert_close(model.state_dict(), state_dict)
        else:
            with pytest.raises(AssertionError):
                torch.testing.assert_close(model.state_dict(), state_dict)

        # No sharding - sharded state dict
        options = StateDictOptions(full_state_dict=False)
        with pytest.raises(ValueError, match="cannot be used"):
            load_model_state_dict(state_dict, model, options, self.rank)

        # Sharding - full state dict
        for kwargs in ({"cpu_offload": True}, {"cpu_offload": False}, {"rank0_only": True}):
            model = MyModel(4).to(device=device)
            sharded_model = thunder.distributed.fsdp(model)
            options = StateDictOptions(full_state_dict=True, **kwargs)
            load_model_state_dict(state_dict, sharded_model, options, self.rank)
            _unshard_params(sharded_model, pg)
            torch.testing.assert_close(model.state_dict(), state_dict)

        # Create a sharded state_dict that can be loaded
        model = MyModel(4).to(device=device)
        sharded_model_expected = thunder.distributed.fsdp(model)
        options = StateDictOptions(full_state_dict=False)
        sharded_state_dict = get_model_state_dict(sharded_model_expected, options, self.rank)

        # Sharding - sharded state dict
        for kwargs in ({"cpu_offload": True}, {"cpu_offload": False}, {"rank0_only": True}):
            model = MyModel(4).to(device=device)
            sharded_model = thunder.distributed.fsdp(model)
            options = StateDictOptions(full_state_dict=False, **kwargs)
            load_model_state_dict(sharded_state_dict, sharded_model, options, self.rank)
            torch.testing.assert_close(sharded_model.state_dict(), sharded_model_expected.state_dict())
//  [python]
def test_quantile_quantization():
    for i in range(100):
        A1 = torch.randn(1024, 1024, device="cuda")
        code = F.estimate_quantiles(A1)
        C = F.quantize_no_absmax(A1, code)
        A2 = F.dequantize_no_absmax(C, code)
        diff = torch.abs(A1 - A2).mean().item()
        assert diff < 0.0075

        A1 = torch.rand(1024, 1024, device="cuda")
        code = F.estimate_quantiles(A1)
        C = F.quantize_no_absmax(A1, code)
        A2 = F.dequantize_no_absmax(C, code)
        diff = torch.abs(A1 - A2).mean().item()
        torch.testing.assert_allclose(A1, A2, atol=5e-3, rtol=0)
        assert diff < 0.001
//  [python]
def test_quantile_quantization():
    for i in range(100):
        A1 = torch.randn(1024, 1024, device="cuda")
        code = F.estimate_quantiles(A1)
        C = F.quantize_no_absmax(A1, code)
        A2 = F.dequantize_no_absmax(C, code)
        diff = torch.abs(A1 - A2).mean().item()
        assert diff < 0.0075

        A1 = torch.rand(1024, 1024, device="cuda")
        code = F.estimate_quantiles(A1)
        C = F.quantize_no_absmax(A1, code)
        A2 = F.dequantize_no_absmax(C, code)
        diff = torch.abs(A1 - A2).mean().item()
        torch.testing.assert_close(A1, A2, atol=5e-3, rtol=0)
        assert diff < 0.001