jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_fixed_size_experience_split_strategy(self):
        x = torch.rand(32, 10)
        y = torch.arange(
            32
        )  # we use ordered labels to reconstruct the order after shuffling
        dd = AvalancheDataset([TensorDataset(x, y)])

        d1, d2 = split_validation_random(validation_size=10, shuffle=True, dataset=dd)
        assert len(d1) + len(d2) == len(dd)

        # check data is shuffled
        iis = []
        for x, ii in d1:
            iis.append(ii)
        for x, ii in d2:
            iis.append(ii)
        assert not all(b >= a for a, b in zip(iis[:-1], iis[1:]))

        # check d1
        for x, ii in d1:
            torch.testing.assert_allclose(x, dd[ii.item()][0])
            torch.testing.assert_allclose(ii, dd[ii.item()][1])
        # check d2
        for x, ii in d2:
            torch.testing.assert_allclose(x, dd[ii.item()][0])
            torch.testing.assert_allclose(ii, dd[ii.item()][1])
//  [python]
def test_fixed_size_experience_split_strategy(self):
        x = torch.rand(32, 10)
        y = torch.arange(
            32
        )  # we use ordered labels to reconstruct the order after shuffling
        dd = AvalancheDataset([TensorDataset(x, y)])

        d1, d2 = split_validation_random(validation_size=10, shuffle=True, dataset=dd)
        assert len(d1) + len(d2) == len(dd)

        # check data is shuffled
        iis = []
        for x, ii in d1:
            iis.append(ii)
        for x, ii in d2:
            iis.append(ii)
        assert not all(b >= a for a, b in zip(iis[:-1], iis[1:]))

        # check d1
        for x, ii in d1:
            torch.testing.assert_close(x, dd[ii.item()][0])
            torch.testing.assert_close(ii, dd[ii.item()][1])
        # check d2
        for x, ii in d2:
            torch.testing.assert_close(x, dd[ii.item()][0])
            torch.testing.assert_close(ii, dd[ii.item()][1])
//  [python]
def test_cg_fixed(self,):
        A, b, x = self.sample()
        cg = CG()
        t.assert_close(x, cg(A, b), atol=1e-5, rtol=1e-4)
//  [python]
def test_cg_fixed(self,):
        A, b, x = self.sample()
        cg = CG()
        torch.testing.assert_close(x, cg(A, b), atol=1e-5, rtol=1e-4)
//  [python]
def test_label_smoothing_algorithm(self, tensors, smoothing, target_type, empty_logger, minimal_state):
        (outputs, target_indices, target_onehot) = tensors

        target = target_indices if target_type == 'indices' else target_onehot

        algorithm = LabelSmoothing(smoothing=smoothing)
        state = minimal_state
        state.batch = (torch.Tensor(), target)
        state.outputs = outputs

        # BEFORE_LOSS should smooth the labels
        algorithm.apply(Event.BEFORE_LOSS, state, empty_logger)
        smoothed_reference = self.reference_smooth_labels(target_onehot, smoothing)

        _, labels = state.batch
        torch.testing.assert_allclose(labels, smoothed_reference)

        # AFTER_LOSS should restore the original targets
        algorithm.apply(Event.AFTER_LOSS, state, empty_logger)

        _, labels = state.batch
        torch.testing.assert_allclose(labels, target)
//  [python]
def test_label_smoothing_algorithm(self, tensors, smoothing, target_type, empty_logger, minimal_state):
        (outputs, target_indices, target_onehot) = tensors

        target = target_indices if target_type == 'indices' else target_onehot

        algorithm = LabelSmoothing(smoothing=smoothing)
        state = minimal_state
        state.batch = (torch.Tensor(), target)
        state.outputs = outputs

        # BEFORE_LOSS should smooth the labels
        algorithm.apply(Event.BEFORE_LOSS, state, empty_logger)
        smoothed_reference = self.reference_smooth_labels(target_onehot, smoothing)

        _, labels = state.batch
        torch.testing.assert_close(labels, smoothed_reference)

        # AFTER_LOSS should restore the original targets
        algorithm.apply(Event.AFTER_LOSS, state, empty_logger)

        _, labels = state.batch
        torch.testing.assert_close(labels, target)
//  [python]
def test_cg_batch(self,):
        A, b, x = self.sample()
        cg = CG()
        A_batch = torch.stack([A, A])
        b_batch = torch.stack([b, b])
        x_batch = torch.stack([x, x])
        t.assert_allclose(x_batch, cg(A_batch, b_batch), atol=1e-5, rtol=1e-4)
//  [python]
def test_cg_batch(self,):
        A, b, x = self.sample()
        cg = CG()
        A_batch = torch.stack([A, A])
        b_batch = torch.stack([b, b])
        x_batch = torch.stack([x, x])
        torch.testing.assert_close(x_batch, cg(A_batch, b_batch), atol=1e-5, rtol=1e-4)
//  [python]
def test_cg_csr(self,):
        A, b, x = self.sample()
        cg = CG()
        A_csr = A.to_sparse_csr()
        t.assert_close(x, cg(A_csr, b), atol=1e-5, rtol=1e-4)
//  [python]
def test_cg_csr(self,):
        A, b, x = self.sample()
        cg = CG()
        A_csr = A.to_sparse_csr()
        torch.testing.assert_close(x, cg(A_csr, b), atol=1e-5, rtol=1e-4)
//  [python]
def test_paged_attention(self, conti_q, blocked_kv, block_offsets,
                             start_loc, seq_lens, history_lens, conti_gt):
        from lmdeploy.pytorch.kernels import paged_attention_fwd
        kv_seq_lens = seq_lens + history_lens
        max_seq_len = seq_lens.max().item()

        blocked_k, blocked_v = blocked_kv
        out = torch.empty_like(conti_q)

        paged_attention_fwd(conti_q,
                            blocked_k,
                            blocked_v,
                            out,
                            block_offsets=block_offsets,
                            q_start_loc=start_loc,
                            q_seqlens=seq_lens,
                            kv_seqlens=kv_seq_lens,
                            max_seqlen=max_seq_len)
        torch.testing.assert_allclose(out, conti_gt, atol=1e-3, rtol=1e-5)
//  [python]
def test_paged_attention(self, conti_q, blocked_kv, block_offsets,
                             start_loc, seq_lens, history_lens, conti_gt):
        from lmdeploy.pytorch.kernels import paged_attention_fwd
        kv_seq_lens = seq_lens + history_lens
        max_seq_len = seq_lens.max().item()

        blocked_k, blocked_v = blocked_kv
        out = torch.empty_like(conti_q)

        paged_attention_fwd(conti_q,
                            blocked_k,
                            blocked_v,
                            out,
                            block_offsets=block_offsets,
                            q_start_loc=start_loc,
                            q_seqlens=seq_lens,
                            kv_seqlens=kv_seq_lens,
                            max_seqlen=max_seq_len)
        torch.testing.assert_close(out, conti_gt, atol=1e-3, rtol=1e-5)
//  [python]
def test_window_attention(self, conti_q, blocked_kv, block_offsets,
                              start_loc, seq_lens, history_lens, win_size,
                              window_gt):
        from lmdeploy.pytorch.kernels import paged_attention_fwd
        kv_seq_lens = seq_lens + history_lens
        max_seq_len = seq_lens.max().item()

        blocked_k, blocked_v = blocked_kv
        out = torch.empty_like(conti_q)
        paged_attention_fwd(conti_q,
                            blocked_k,
                            blocked_v,
                            out,
                            block_offsets=block_offsets,
                            q_start_loc=start_loc,
                            q_seqlens=seq_lens,
                            kv_seqlens=kv_seq_lens,
                            max_seqlen=max_seq_len,
                            window_size=win_size)
        torch.allclose(out, window_gt, atol=1e-3, rtol=1e-5)
//  [python]
def test_window_attention(self, conti_q, blocked_kv, block_offsets,
                              start_loc, seq_lens, history_lens, win_size,
                              window_gt):
        from lmdeploy.pytorch.kernels import paged_attention_fwd
        kv_seq_lens = seq_lens + history_lens
        max_seq_len = seq_lens.max().item()

        blocked_k, blocked_v = blocked_kv
        out = torch.empty_like(conti_q)
        paged_attention_fwd(conti_q,
                            blocked_k,
                            blocked_v,
                            out,
                            block_offsets=block_offsets,
                            q_start_loc=start_loc,
                            q_seqlens=seq_lens,
                            kv_seqlens=kv_seq_lens,
                            max_seqlen=max_seq_len,
                            window_size=win_size)
        torch.testing.assert_close(out, window_gt, atol=1e-3, rtol=1e-5)
//  [python]
def assert_most_approx_close(a, b, rtol=1e-3, atol=1e-3, max_error_count=0):
    idx = torch.isclose(a, b, rtol=rtol, atol=atol)
    error_count = (idx == 0).sum().item()
    if error_count > max_error_count:
        print(f"Too many values not close: assert {error_count} < {max_error_count}")
        t.assert_close(a, b, rtol=rtol, atol=atol)
//  [python]
def assert_most_approx_close(a, b, rtol=1e-3, atol=1e-3, max_error_count=0):
    idx = torch.isclose(a, b, rtol=rtol, atol=atol)
    error_count = (idx == 0).sum().item()
    if error_count > max_error_count:
        print(f"Too many values not close: assert {error_count} < {max_error_count}")
        torch.testing.assert_close(a, b, rtol=rtol, atol=atol)
//  [python]
def test_equal_dict(first: Dict, second: Dict, sub_paths: Optional[List[str]] = None) -> None:
    """Raise if doesn't match."""

    if sub_paths is None:
        sub_paths = []

    first_keys = set(first.keys())
    second_keys = set(second.keys())
    assert first_keys == second_keys, f"Keys don't match.\nFirst: {first_keys}\nSecond: {second_keys}"
    for key in first_keys:
        first_elt = first[key]
        second_elt = second[key]

        if isinstance(first_elt, dict):
            assert isinstance(second_elt, dict), f"{first_elt} doesn't match {second_elt}"
            test_equal_dict(first_elt, second_elt, sub_paths=sub_paths + [str(key)])
        elif isinstance(first_elt, torch.Tensor):
            assert isinstance(second_elt, torch.Tensor), f"{first_elt} doesn't match {second_elt}"
            torch.testing.assert_allclose(
                first_elt,
                second_elt,
                atol=0.0,
                rtol=0.0,
                msg=lambda msg: f"tensor at {'.'.join(sub_paths + [str(key)])} don't match.\nCur: {first_elt}\nRef: {second_elt}\n{msg}",
            )
        else:
            assert (
                first_elt == second_elt
            ), f"{first_elt} doesn't match {second_elt} at key {'.'.join(sub_paths + [str(key)])}"
//  [python]
def test_equal_dict(first: Dict, second: Dict, sub_paths: Optional[List[str]] = None) -> None:
    """Raise if doesn't match."""

    if sub_paths is None:
        sub_paths = []

    first_keys = set(first.keys())
    second_keys = set(second.keys())
    assert first_keys == second_keys, f"Keys don't match.\nFirst: {first_keys}\nSecond: {second_keys}"
    for key in first_keys:
        first_elt = first[key]
        second_elt = second[key]

        if isinstance(first_elt, dict):
            assert isinstance(second_elt, dict), f"{first_elt} doesn't match {second_elt}"
            test_equal_dict(first_elt, second_elt, sub_paths=sub_paths + [str(key)])
        elif isinstance(first_elt, torch.Tensor):
            assert isinstance(second_elt, torch.Tensor), f"{first_elt} doesn't match {second_elt}"
            torch.testing.assert_close(
                first_elt,
                second_elt,
                atol=0.0,
                rtol=0.0,
                msg=lambda msg: f"tensor at {'.'.join(sub_paths + [str(key)])} don't match.\nCur: {first_elt}\nRef: {second_elt}\n{msg}",
            )
        else:
            assert (
                first_elt == second_elt
            ), f"{first_elt} doesn't match {second_elt} at key {'.'.join(sub_paths + [str(key)])}"
//  [python]
def test_flash_attention_stride0(self):
        input_shape = (
            1,
            16,
            1,
            48,
        )
        input_stride = (
            0,
            48,
            0,
            1,
        )
        q = torch.randn(
            input_shape, device="cpu", dtype=torch.float32, requires_grad=False
        ).as_strided(input_shape, input_stride)
        k = torch.randn(
            input_shape, device="cpu", dtype=torch.float32, requires_grad=False
        ).as_strided(input_shape, input_stride)
        v = torch.randn(
            input_shape, device="cpu", dtype=torch.float32, requires_grad=False
        ).as_strided(input_shape, input_stride)
        atol = 1e-5
        rtol = 5e-6
        q2 = q.clone()
        k2 = k.clone()
        v2 = v.clone()
        actual = torch.ops.torch_ipex.flash_attention(q, k, v)[0]
        math_ref = torch._scaled_dot_product_attention_math(q2, k2, v2)[0]
        torch.testing.assert_allclose(actual, math_ref, atol=1e-5, rtol=5e-6)
//  [python]
def test_flash_attention_stride0(self):
        input_shape = (
            1,
            16,
            1,
            48,
        )
        input_stride = (
            0,
            48,
            0,
            1,
        )
        q = torch.randn(
            input_shape, device="cpu", dtype=torch.float32, requires_grad=False
        ).as_strided(input_shape, input_stride)
        k = torch.randn(
            input_shape, device="cpu", dtype=torch.float32, requires_grad=False
        ).as_strided(input_shape, input_stride)
        v = torch.randn(
            input_shape, device="cpu", dtype=torch.float32, requires_grad=False
        ).as_strided(input_shape, input_stride)
        atol = 1e-5
        rtol = 5e-6
        q2 = q.clone()
        k2 = k.clone()
        v2 = v.clone()
        actual = torch.ops.torch_ipex.flash_attention(q, k, v)[0]
        math_ref = torch._scaled_dot_product_attention_math(q2, k2, v2)[0]
        torch.testing.assert_close(actual, math_ref, atol=1e-5, rtol=5e-6)
//  [python]
def test_permutation_autograd(self) -> None:
        net = Net().to(self.device)

        input = torch.randn(2, 1).to(self.device)
        input_sum = input.sum().item()

        output = net(input)
        output.sum().backward()

        # check grads for fc1 when permuted, equals to fc2 weights times input_sum
        # pyre-fixme[16]: Optional type has no attribute `view`.
        permute_res = net.permute_pooled_embeddings(net.fc1.weight.grad.view(1, 10))
        permute_ref = input_sum * net.fc2.weight
        self.assertEqual(permute_res, permute_ref, rtol=1e-03, atol=1e-03)
//  [python]
def test_permutation_autograd(self) -> None:
        net = Net().to(self.device)

        input = torch.randn(2, 1).to(self.device)
        input_sum = input.sum().item()

        output = net(input)
        output.sum().backward()

        # check grads for fc1 when permuted, equals to fc2 weights times input_sum
        # pyre-fixme[16]: Optional type has no attribute `view`.
        permute_res = net.permute_pooled_embeddings(net.fc1.weight.grad.view(1, 10))
        permute_ref = input_sum * net.fc2.weight
        torch.testing.assert_close(permute_res, permute_ref, rtol=1e-03, atol=1e-03)
//  [python]
def test_bert(batch_size: int, seq_length: int, use_fp16, use_tensor_core, dynamic):
    tokens_tensor = torch.zeros((batch_size, seq_length), dtype=torch.long, device='cuda')
    segments_tensors = torch.zeros((batch_size, seq_length), dtype=torch.long, device='cuda')
    args = (tokens_tensor.cuda(),)
    kwargs = {'token_type_ids': segments_tensors.cuda()}
    model = torch.hub.load('huggingface/pytorch-transformers', 'model', 'bert-base-uncased').cuda().eval()
    model_opt = torch.compile(model, backend='hidet', dynamic=dynamic)
    y1 = model(*args, **kwargs).last_hidden_state

    try:
        hidet.torch.dynamo_config.use_fp16(use_fp16)
        hidet.torch.dynamo_config.use_tensor_core(use_tensor_core)

        y2 = model_opt(*args, **kwargs).last_hidden_state
        tol = 1e-1 if use_fp16 else 1e-2
        torch.testing.assert_allclose(y1, y2, atol=tol, rtol=tol)
    finally:
        # in case of failure, reset the config
        hidet.torch.dynamo_config.reset()
//  [python]
def test_bert(batch_size: int, seq_length: int, use_fp16, use_tensor_core, dynamic):
    tokens_tensor = torch.zeros((batch_size, seq_length), dtype=torch.long, device='cuda')
    segments_tensors = torch.zeros((batch_size, seq_length), dtype=torch.long, device='cuda')
    args = (tokens_tensor.cuda(),)
    kwargs = {'token_type_ids': segments_tensors.cuda()}
    model = torch.hub.load('huggingface/pytorch-transformers', 'model', 'bert-base-uncased').cuda().eval()
    model_opt = torch.compile(model, backend='hidet', dynamic=dynamic)
    y1 = model(*args, **kwargs).last_hidden_state

    try:
        hidet.torch.dynamo_config.use_fp16(use_fp16)
        hidet.torch.dynamo_config.use_tensor_core(use_tensor_core)

        y2 = model_opt(*args, **kwargs).last_hidden_state
        tol = 1e-1 if use_fp16 else 1e-2
        torch.testing.assert_close(y1, y2, atol=tol, rtol=tol)
    finally:
        # in case of failure, reset the config
        hidet.torch.dynamo_config.reset()
//  [python]
def test_icdf(self):
        params = TensorDict(
            {
                "cont": {
                    "loc": torch.randn(3, 4, requires_grad=True),
                    "scale": torch.rand(3, 4, requires_grad=True),
                },
                ("nested", "cont"): {
                    "loc": torch.randn(3, 4, requires_grad=True),
                    "scale": torch.rand(3, 4, requires_grad=True),
                },
            },
            [3],
        )
        dist = CompositeDistribution(
            params,
            distribution_map={
                "cont": distributions.Normal,
                ("nested", "cont"): distributions.Normal,
            },
        )
        sample = dist.rsample((4,))
        sample = dist.cdf(sample)
        sample = dist.icdf(sample)
        assert sample.get("cont_icdf").requires_grad
        assert sample.get(("nested", "cont_icdf")).requires_grad
        torch.testing.assert_allclose(sample.get("cont"), sample.get("cont_icdf"))
//  [python]
def test_icdf(self):
        params = TensorDict(
            {
                "cont": {
                    "loc": torch.randn(3, 4, requires_grad=True),
                    "scale": torch.rand(3, 4, requires_grad=True),
                },
                ("nested", "cont"): {
                    "loc": torch.randn(3, 4, requires_grad=True),
                    "scale": torch.rand(3, 4, requires_grad=True),
                },
            },
            [3],
        )
        dist = CompositeDistribution(
            params,
            distribution_map={
                "cont": distributions.Normal,
                ("nested", "cont"): distributions.Normal,
            },
        )
        sample = dist.rsample((4,))
        sample = dist.cdf(sample)
        sample = dist.icdf(sample)
        assert sample.get("cont_icdf").requires_grad
        assert sample.get(("nested", "cont_icdf")).requires_grad
        torch.testing.assert_close(sample.get("cont"), sample.get("cont_icdf"))
//  [python]
def test_prob_module(self, interaction, return_log_prob):
        params = TensorDict(
            {
                "params": {
                    "cont": {
                        "loc": torch.randn(3, 4, requires_grad=True),
                        "scale": torch.rand(3, 4, requires_grad=True),
                    },
                    ("nested", "cont"): {
                        "loc": torch.randn(3, 4, requires_grad=True),
                        "scale": torch.rand(3, 4, requires_grad=True),
                    },
                }
            },
            [3],
        )
        in_keys = ["params"]
        out_keys = ["cont", ("nested", "cont")]
        distribution_map = {
            "cont": distributions.Normal,
            ("nested", "cont"): distributions.Normal,
        }
        module = ProbabilisticTensorDictModule(
            in_keys=in_keys,
            out_keys=out_keys,
            distribution_class=CompositeDistribution,
            distribution_kwargs={"distribution_map": distribution_map},
            default_interaction_type=interaction,
            return_log_prob=return_log_prob,
        )
        sample = module(params)
        if return_log_prob:
            assert "cont_log_prob" in sample.keys()
            assert ("nested", "cont_log_prob") in sample.keys(True)
        sample_clone = sample.clone()
        lp = module.log_prob(sample_clone)
        if return_log_prob:
            t.assert_close(
                lp,
                sample.get("cont_log_prob").sum(-1)
                + sample.get(("nested", "cont_log_prob")).sum(-1),
            )
        else:
            t.assert_close(
                lp,
                sample_clone.get("cont_log_prob").sum(-1)
                + sample_clone.get(("nested", "cont_log_prob")).sum(-1),
            )
//  [python]
def test_prob_module(self, interaction, return_log_prob):
        params = TensorDict(
            {
                "params": {
                    "cont": {
                        "loc": torch.randn(3, 4, requires_grad=True),
                        "scale": torch.rand(3, 4, requires_grad=True),
                    },
                    ("nested", "cont"): {
                        "loc": torch.randn(3, 4, requires_grad=True),
                        "scale": torch.rand(3, 4, requires_grad=True),
                    },
                }
            },
            [3],
        )
        in_keys = ["params"]
        out_keys = ["cont", ("nested", "cont")]
        distribution_map = {
            "cont": distributions.Normal,
            ("nested", "cont"): distributions.Normal,
        }
        module = ProbabilisticTensorDictModule(
            in_keys=in_keys,
            out_keys=out_keys,
            distribution_class=CompositeDistribution,
            distribution_kwargs={"distribution_map": distribution_map},
            default_interaction_type=interaction,
            return_log_prob=return_log_prob,
        )
        sample = module(params)
        if return_log_prob:
            assert "cont_log_prob" in sample.keys()
            assert ("nested", "cont_log_prob") in sample.keys(True)
        sample_clone = sample.clone()
        lp = module.log_prob(sample_clone)
        if return_log_prob:
            torch.testing.assert_close(
                lp,
                sample.get("cont_log_prob").sum(-1)
                + sample.get(("nested", "cont_log_prob")).sum(-1),
            )
        else:
            torch.testing.assert_close(
                lp,
                sample_clone.get("cont_log_prob").sum(-1)
                + sample_clone.get(("nested", "cont_log_prob")).sum(-1),
            )
//  [python]
def test_optimizer32bit(dim1, dim2, gtype, optim_name):
    if gtype == torch.bfloat16 and optim_name in ["momentum", "rmsprop"]:
        pytest.skip()
    if dim1 == 1 and dim2 == 1:
        return
    p1 = torch.randn(dim1, dim2, device="cuda", dtype=gtype) * 0.1
    p2 = p1.clone()
    p1 = p1.float()

    torch_optimizer = str2optimizers[optim_name][0]([p1])
    bnb_optimizer = str2optimizers[optim_name][1]([p2])

    if gtype == torch.float32:
        atol, rtol = 1e-6, 1e-5
    elif gtype == torch.bfloat16:
        atol, rtol = 1e-3, 1e-2
    else:
        atol, rtol = 1e-4, 1e-3

    for i in range(k):
        g = torch.randn(dim1, dim2, device="cuda", dtype=gtype) * 0.01
        p1.grad = g.clone().float()
        p2.grad = g.clone()

        bnb_optimizer.step()
        torch_optimizer.step()

        for name1, name2 in str2statenames[optim_name]:
            t.assert_close(
                torch_optimizer.state[p1][name1],
                bnb_optimizer.state[p2][name2].cuda(),
                atol=atol,
                rtol=rtol,
            )

        # since Lion can have pretty noisy updates where things lie at the boundary
        # allow up to 10 errors for Lion
        assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=10)

        if i % (k // 5) == 0 and i > 0:
            path = get_temp_dir()
            torch.save(bnb_optimizer.state_dict(), join(path, "opt.pt"))
            del bnb_optimizer
            bnb_optimizer = None
            bnb_optimizer = str2optimizers[optim_name][1]([p2])
            bnb_optimizer.load_state_dict(torch.load(join(path, "opt.pt")))
            rm_path(path)
            # since Lion can have pretty noisy updates where things lie at the boundary
            # allow up to 10 errors for Lion
            assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=10)
            for name1, name2 in str2statenames[optim_name]:
                # since Lion can have pretty noisy updates where things lie at the boundary
                # allow up to 10 errors for Lion
                assert_most_approx_close(
                    torch_optimizer.state[p1][name1],
                    bnb_optimizer.state[p2][name2],
                    atol=atol,
                    rtol=rtol,
                    max_error_count=10,
                )

        if gtype != torch.float32:
            # the adam buffers should also be close because they are 32-bit
            # but the parameters can diverge because they are 16-bit
            # the difference grow larger and larger with each update
            # --> copy the state to keep weights close
            p1.data = p1.data.to(p2.dtype).float()
            p2.copy_(p1.data)
            t.assert_close(p1.to(p2.dtype), p2)
        if optim_name in ["lars", "lamb"]:
            assert bnb_optimizer.state[p2]["unorm_vec"] > 0.0
//  [python]
def test_optimizer32bit(dim1, dim2, gtype, optim_name):
    if gtype == torch.bfloat16 and optim_name in ["momentum", "rmsprop"]:
        pytest.skip()
    if dim1 == 1 and dim2 == 1:
        return
    p1 = torch.randn(dim1, dim2, device="cuda", dtype=gtype) * 0.1
    p2 = p1.clone()
    p1 = p1.float()

    torch_optimizer = str2optimizers[optim_name][0]([p1])
    bnb_optimizer = str2optimizers[optim_name][1]([p2])

    if gtype == torch.float32:
        atol, rtol = 1e-6, 1e-5
    elif gtype == torch.bfloat16:
        atol, rtol = 1e-3, 1e-2
    else:
        atol, rtol = 1e-4, 1e-3

    for i in range(k):
        g = torch.randn(dim1, dim2, device="cuda", dtype=gtype) * 0.01
        p1.grad = g.clone().float()
        p2.grad = g.clone()

        bnb_optimizer.step()
        torch_optimizer.step()

        for name1, name2 in str2statenames[optim_name]:
            torch.testing.assert_close(
                torch_optimizer.state[p1][name1],
                bnb_optimizer.state[p2][name2].cuda(),
                atol=atol,
                rtol=rtol,
            )

        # since Lion can have pretty noisy updates where things lie at the boundary
        # allow up to 10 errors for Lion
        assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=10)

        if i % (k // 5) == 0 and i > 0:
            path = get_temp_dir()
            torch.save(bnb_optimizer.state_dict(), join(path, "opt.pt"))
            del bnb_optimizer
            bnb_optimizer = None
            bnb_optimizer = str2optimizers[optim_name][1]([p2])
            bnb_optimizer.load_state_dict(torch.load(join(path, "opt.pt")))
            rm_path(path)
            # since Lion can have pretty noisy updates where things lie at the boundary
            # allow up to 10 errors for Lion
            assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=10)
            for name1, name2 in str2statenames[optim_name]:
                # since Lion can have pretty noisy updates where things lie at the boundary
                # allow up to 10 errors for Lion
                assert_most_approx_close(
                    torch_optimizer.state[p1][name1],
                    bnb_optimizer.state[p2][name2],
                    atol=atol,
                    rtol=rtol,
                    max_error_count=10,
                )

        if gtype != torch.float32:
            # the adam buffers should also be close because they are 32-bit
            # but the parameters can diverge because they are 16-bit
            # the difference grow larger and larger with each update
            # --> copy the state to keep weights close
            p1.data = p1.data.to(p2.dtype).float()
            p2.copy_(p1.data)
            torch.testing.assert_close(p1.to(p2.dtype), p2)
        if optim_name in ["lars", "lamb"]:
            assert bnb_optimizer.state[p2]["unorm_vec"] > 0.0
//  [python]
def test_cm(self, module_name, input_name, as_module, inplace):
        torch.manual_seed(0)
        module = getattr(self, module_name)
        x = getattr(self, input_name)
        params = TensorDict.from_module(module, as_module=as_module)
        params0 = params.clone().zero_()
        y = module(*x)
        with params0.to_module(module, inplace=inplace):
            y0 = module(*x)
            if as_module:
                # if as_module=False, params0 is not made of parameters anymore
                assert (params0 == TensorDict.from_module(module)).all()

                # check identities
                for p1, p2 in zip(
                    TensorDict.from_module(module).values(True, True),
                    params0.values(True, True),
                ):
                    if inplace:
                        assert p1 is not p2
                    else:
                        assert p1 is p2

        y1 = module(*x)
        t.assert_close(y, y1)
        assert (y0 == 0).all()
        assert (y0 != y1).all()
        assert (TensorDict.from_module(module) == params).all()
//  [python]
def test_cm(self, module_name, input_name, as_module, inplace):
        torch.manual_seed(0)
        module = getattr(self, module_name)
        x = getattr(self, input_name)
        params = TensorDict.from_module(module, as_module=as_module)
        params0 = params.clone().zero_()
        y = module(*x)
        with params0.to_module(module, inplace=inplace):
            y0 = module(*x)
            if as_module:
                # if as_module=False, params0 is not made of parameters anymore
                assert (params0 == TensorDict.from_module(module)).all()

                # check identities
                for p1, p2 in zip(
                    TensorDict.from_module(module).values(True, True),
                    params0.values(True, True),
                ):
                    if inplace:
                        assert p1 is not p2
                    else:
                        assert p1 is p2

        y1 = module(*x)
        torch.testing.assert_close(y, y1)
        assert (y0 == 0).all()
        assert (y0 != y1).all()
        assert (TensorDict.from_module(module) == params).all()
//  [python]
def test_adam_percentile_clipping(dim1, dim2, gtype, optim_bits):
    if dim1 == 1 and dim2 == 1:
        return
    p1 = torch.randn(dim1, dim2, device="cpu", dtype=gtype) * 0.1
    beta1 = 0.9
    beta2 = 0.999
    lr = 0.001
    eps = 1e-8
    p1 = p1.cuda()
    p2 = p1.clone()
    adam1 = bnb.optim.Adam([p1], lr, (beta1, beta2), eps, optim_bits=optim_bits)
    adam2 = bnb.optim.Adam(
        [p2],
        lr,
        (beta1, beta2),
        eps,
        optim_bits=optim_bits,
        percentile_clipping=5,
    )

    gnorm_vec = torch.zeros(100).cuda()
    step = 0

    for i in range(50):
        step += 1
        g1 = torch.randn(dim1, dim2, device="cuda", dtype=gtype) * 0.1 + (0.01 * i)
        g2 = g1.clone()
        p2.grad = g2

        current_gnorm, clip_val, gnorm_scale = F.percentile_clipping(g1, gnorm_vec, step, 5)
        g1 = (g1.float() * gnorm_scale).to(gtype)
        p1.grad = g1

        adam1.step()
        adam2.step()

        # gnorm_scale is not deterministic (warp reductions), as such there can be slight differences in state
        if optim_bits == 32:
            torch.testing.assert_allclose(p1, p2)
            torch.testing.assert_allclose(
                adam1.state[p1]["state1"],
                adam2.state[p2]["state1"],
                atol=5e-5,
                rtol=1e-4,
            )
            torch.testing.assert_allclose(
                adam1.state[p1]["state2"],
                adam2.state[p2]["state2"],
                atol=5e-5,
                rtol=1e-4,
            )
        elif optim_bits == 8:
            torch.testing.assert_allclose(p1, p2, atol=1e-4, rtol=1e-3)
            torch.testing.assert_allclose(
                adam1.state[p1]["state1"],
                adam2.state[p2]["state1"],
                atol=2,
                rtol=1e-3,
            )
            torch.testing.assert_allclose(
                adam1.state[p1]["state2"],
                adam2.state[p2]["state2"],
                atol=2,
                rtol=1e-3,
            )
            adam1.state[p1]["state1"].copy_(adam2.state[p2]["state1"])
            adam1.state[p1]["state2"].copy_(adam2.state[p2]["state2"])
        if i % 10 == 0 and i > 0:
            path = get_temp_dir()
            torch.save(adam2.state_dict(), join(path, "opt.pt"))
            del adam2
            adam2 = None
            adam2 = bnb.optim.Adam(
                [p2],
                lr,
                (beta1, beta2),
                eps,
                optim_bits=optim_bits,
                percentile_clipping=5,
            )
            adam2.load_state_dict(torch.load(join(path, "opt.pt")))
//  [python]
def test_adam_percentile_clipping(dim1, dim2, gtype, optim_bits):
    if dim1 == 1 and dim2 == 1:
        return
    p1 = torch.randn(dim1, dim2, device="cpu", dtype=gtype) * 0.1
    beta1 = 0.9
    beta2 = 0.999
    lr = 0.001
    eps = 1e-8
    p1 = p1.cuda()
    p2 = p1.clone()
    adam1 = bnb.optim.Adam([p1], lr, (beta1, beta2), eps, optim_bits=optim_bits)
    adam2 = bnb.optim.Adam(
        [p2],
        lr,
        (beta1, beta2),
        eps,
        optim_bits=optim_bits,
        percentile_clipping=5,
    )

    gnorm_vec = torch.zeros(100).cuda()
    step = 0

    for i in range(50):
        step += 1
        g1 = torch.randn(dim1, dim2, device="cuda", dtype=gtype) * 0.1 + (0.01 * i)
        g2 = g1.clone()
        p2.grad = g2

        current_gnorm, clip_val, gnorm_scale = F.percentile_clipping(g1, gnorm_vec, step, 5)
        g1 = (g1.float() * gnorm_scale).to(gtype)
        p1.grad = g1

        adam1.step()
        adam2.step()

        # gnorm_scale is not deterministic (warp reductions), as such there can be slight differences in state
        if optim_bits == 32:
            torch.testing.assert_close(p1, p2)
            torch.testing.assert_close(
                adam1.state[p1]["state1"],
                adam2.state[p2]["state1"],
                atol=5e-5,
                rtol=1e-4,
            )
            torch.testing.assert_close(
                adam1.state[p1]["state2"],
                adam2.state[p2]["state2"],
                atol=5e-5,
                rtol=1e-4,
            )
        elif optim_bits == 8:
            torch.testing.assert_close(p1, p2, atol=1e-4, rtol=1e-3)
            torch.testing.assert_close(
                adam1.state[p1]["state1"],
                adam2.state[p2]["state1"],
                atol=2,
                rtol=1e-3,
            )
            torch.testing.assert_close(
                adam1.state[p1]["state2"],
                adam2.state[p2]["state2"],
                atol=2,
                rtol=1e-3,
            )
            adam1.state[p1]["state1"].copy_(adam2.state[p2]["state1"])
            adam1.state[p1]["state2"].copy_(adam2.state[p2]["state2"])
        if i % 10 == 0 and i > 0:
            path = get_temp_dir()
            torch.save(adam2.state_dict(), join(path, "opt.pt"))
            del adam2
            adam2 = None
            adam2 = bnb.optim.Adam(
                [p2],
                lr,
                (beta1, beta2),
                eps,
                optim_bits=optim_bits,
                percentile_clipping=5,
            )
            adam2.load_state_dict(torch.load(join(path, "opt.pt")))
//  [python]
def test_static(self, module_name, input_name, as_module, inplace):
        torch.manual_seed(0)
        module = getattr(self, module_name)
        x = getattr(self, input_name)
        params = TensorDict.from_module(module, as_module=as_module)
        if inplace:
            params = params.clone()
        params0 = params.clone().zero_()
        y = module(*x)
        params0.to_module(module, inplace=inplace)
        y0 = module(*x)
        # check identities
        for k, p1, p2 in zip(
            params0.keys(True, True),
            TensorDict.from_module(module).values(True, True),
            params0.values(True, True),
        ):
            if inplace:
                assert p1 is not p2, k
            else:
                assert p1 is p2, k
        params.to_module(module, inplace=inplace)
        # check identities
        for p1, p2 in zip(
            TensorDict.from_module(module).values(True, True),
            params.values(True, True),
        ):
            if inplace:
                assert p1 is not p2
            else:
                assert p1 is p2
        y1 = module(*x)
        torch.testing.assert_allclose(y, y1)
        assert (y0 == 0).all()
        assert (y0 != y1).all()
//  [python]
def test_static(self, module_name, input_name, as_module, inplace):
        torch.manual_seed(0)
        module = getattr(self, module_name)
        x = getattr(self, input_name)
        params = TensorDict.from_module(module, as_module=as_module)
        if inplace:
            params = params.clone()
        params0 = params.clone().zero_()
        y = module(*x)
        params0.to_module(module, inplace=inplace)
        y0 = module(*x)
        # check identities
        for k, p1, p2 in zip(
            params0.keys(True, True),
            TensorDict.from_module(module).values(True, True),
            params0.values(True, True),
        ):
            if inplace:
                assert p1 is not p2, k
            else:
                assert p1 is p2, k
        params.to_module(module, inplace=inplace)
        # check identities
        for p1, p2 in zip(
            TensorDict.from_module(module).values(True, True),
            params.values(True, True),
        ):
            if inplace:
                assert p1 is not p2
            else:
                assert p1 is p2
        y1 = module(*x)
        torch.testing.assert_close(y, y1)
        assert (y0 == 0).all()
        assert (y0 != y1).all()
//  [python]
def test_prob_module_seq(self, interaction, return_log_prob):
        params = TensorDict(
            {
                "params": {
                    "cont": {
                        "loc": torch.randn(3, 4, requires_grad=True),
                        "scale": torch.rand(3, 4, requires_grad=True),
                    },
                    ("nested", "cont"): {
                        "loc": torch.randn(3, 4, requires_grad=True),
                        "scale": torch.rand(3, 4, requires_grad=True),
                    },
                }
            },
            [3],
        )
        in_keys = ["params"]
        out_keys = ["cont", ("nested", "cont")]
        distribution_map = {
            "cont": distributions.Normal,
            ("nested", "cont"): distributions.Normal,
        }
        backbone = TensorDictModule(lambda: None, in_keys=[], out_keys=[])
        module = ProbabilisticTensorDictSequential(
            backbone,
            ProbabilisticTensorDictModule(
                in_keys=in_keys,
                out_keys=out_keys,
                distribution_class=CompositeDistribution,
                distribution_kwargs={"distribution_map": distribution_map},
                default_interaction_type=interaction,
                return_log_prob=return_log_prob,
            ),
        )
        sample = module(params)
        if return_log_prob:
            assert "cont_log_prob" in sample.keys()
            assert ("nested", "cont_log_prob") in sample.keys(True)
        sample_clone = sample.clone()
        lp = module.log_prob(sample_clone)
        if return_log_prob:
            torch.testing.assert_allclose(
                lp,
                sample.get("cont_log_prob").sum(-1)
                + sample.get(("nested", "cont_log_prob")).sum(-1),
            )
        else:
            torch.testing.assert_allclose(
                lp,
                sample_clone.get("cont_log_prob").sum(-1)
                + sample_clone.get(("nested", "cont_log_prob")).sum(-1),
            )
//  [python]
def test_prob_module_seq(self, interaction, return_log_prob):
        params = TensorDict(
            {
                "params": {
                    "cont": {
                        "loc": torch.randn(3, 4, requires_grad=True),
                        "scale": torch.rand(3, 4, requires_grad=True),
                    },
                    ("nested", "cont"): {
                        "loc": torch.randn(3, 4, requires_grad=True),
                        "scale": torch.rand(3, 4, requires_grad=True),
                    },
                }
            },
            [3],
        )
        in_keys = ["params"]
        out_keys = ["cont", ("nested", "cont")]
        distribution_map = {
            "cont": distributions.Normal,
            ("nested", "cont"): distributions.Normal,
        }
        backbone = TensorDictModule(lambda: None, in_keys=[], out_keys=[])
        module = ProbabilisticTensorDictSequential(
            backbone,
            ProbabilisticTensorDictModule(
                in_keys=in_keys,
                out_keys=out_keys,
                distribution_class=CompositeDistribution,
                distribution_kwargs={"distribution_map": distribution_map},
                default_interaction_type=interaction,
                return_log_prob=return_log_prob,
            ),
        )
        sample = module(params)
        if return_log_prob:
            assert "cont_log_prob" in sample.keys()
            assert ("nested", "cont_log_prob") in sample.keys(True)
        sample_clone = sample.clone()
        lp = module.log_prob(sample_clone)
        if return_log_prob:
            torch.testing.assert_close(
                lp,
                sample.get("cont_log_prob").sum(-1)
                + sample.get(("nested", "cont_log_prob")).sum(-1),
            )
        else:
            torch.testing.assert_close(
                lp,
                sample_clone.get("cont_log_prob").sum(-1)
                + sample_clone.get(("nested", "cont_log_prob")).sum(-1),
            )
//  [python]
def run_worker(args, model_class):
    mod = model_class()
    x = torch.randn(batch_size, d_hid)
    y = torch.randn(batch_size, d_hid)

    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(x, y),
    )

    ref_out = mod(x, y)
    out = pipe(x, y)
    t.assert_close(out, ref_out)
    print(f"equivalence test passed loss={out[1]} ref_loss={ref_out[1]}")
//  [python]
def run_worker(args, model_class):
    mod = model_class()
    x = torch.randn(batch_size, d_hid)
    y = torch.randn(batch_size, d_hid)

    pipe = pipeline(
        mod,
        args.chunks,
        example_args=(x, y),
    )

    ref_out = mod(x, y)
    out = pipe(x, y)
    torch.testing.assert_close(out, ref_out)
    print(f"equivalence test passed loss={out[1]} ref_loss={ref_out[1]}")
//  [python]
def test_cm_meta(self, module_name, input_name, as_module):
        torch.manual_seed(0)
        module = getattr(self, module_name)
        x = getattr(self, input_name)
        params = TensorDict.from_module(module, as_module=as_module)
        params_meta = params.detach().to("meta")
        y = module(*x)
        with params_meta.to_module(module):
            module_meta = copy.deepcopy(module)
        y1 = module(*x)
        with params.to_module(module_meta):
            y2 = module_meta(*x)
        t.assert_close(y, y1)
        t.assert_close(y, y2)
        assert (TensorDict.from_module(module) == params).all()
//  [python]
def test_cm_meta(self, module_name, input_name, as_module):
        torch.manual_seed(0)
        module = getattr(self, module_name)
        x = getattr(self, input_name)
        params = TensorDict.from_module(module, as_module=as_module)
        params_meta = params.detach().to("meta")
        y = module(*x)
        with params_meta.to_module(module):
            module_meta = copy.deepcopy(module)
        y1 = module(*x)
        with params.to_module(module_meta):
            y2 = module_meta(*x)
        torch.testing.assert_close(y, y1)
        torch.testing.assert_close(y, y2)
        assert (TensorDict.from_module(module) == params).all()
//  [python]
def _test_against_reference(self, inputs, reference, eps, dtype):
        q, k, v, rpb, kernel_size, dilation, is_causal = inputs
        kernel_size, dilation, is_causal = check_args(kernel_size, dilation, is_causal)
        out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            rpb_ = rpb if rpb is None else rpb.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            out = na3d(
                q_,
                k_,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
                rpb=rpb_,
            )

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

            out = na3d(
                q_,
                k_,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
                rpb=rpb_,
            )

            torch.testing.assert_close(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def _test_against_sdpa(self, inputs, reference, eps, dtype):
        q, k, v = inputs
        kernel_size = (q.shape[1], q.shape[2], q.shape[3])
        kernel_size, dilation, is_causal = check_args(kernel_size, 1, False)
        out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            out = na3d(
                q_,
                k_,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
            )

            torch.testing.assert_allclose(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def _test_against_sdpa(self, inputs, reference, eps, dtype):
        q, k, v = inputs
        kernel_size = (q.shape[1], q.shape[2], q.shape[3])
        kernel_size, dilation, is_causal = check_args(kernel_size, 1, False)
        out_ref = reference
        with torch.no_grad():
            q_, k_, v_ = q.clone().to(dtype), k.clone().to(dtype), v.clone().to(dtype)
            assert q_.is_cuda and k_.is_cuda and v_.is_cuda

            out = na3d(
                q_,
                k_,
                v_,
                kernel_size=kernel_size,
                dilation=dilation,
                is_causal=is_causal,
            )

            torch.testing.assert_close(out.float(), out_ref, atol=eps, rtol=0)
//  [python]
def testScaledDotProductAttention():
    from deepchem.models import ScaledDotProductAttention as SDPA
    attn = SDPA()
    x = torch.ones(2, 5)
    # Linear layers for making query, key, value
    Q, K, V = nn.Parameter(torch.ones(5)), nn.Parameter(
        torch.ones(5)), nn.Parameter(torch.ones(5))
    query, key, value = Q * x, K * x, V * x
    mask = torch.Tensor([1, 0])
    x_out, attn_score = attn(query, key, value, mask=mask)
    torch.testing.assert_allclose(x_out, torch.ones(2, 5))
    torch.testing.assert_allclose(attn_score, torch.Tensor([[1, 0], [1, 0]]))