jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_prob(self):
        x = torch.tensor([[[0.3, 0.6, 0.1]], [[0.6, 0.3, 0.2]]]).log()
        decoder = CTCGreedyDecoder()
        r = decoder(x, segmentation=True, apply_softmax=False)
        e = [[1]]
        self.assertEqual(e, r["hyp"])
        # Check actual loss prob
        paths = torch.tensor(
            [x[0, 0, a] + x[1, 0, b] for a, b in ((0, 1), (1, 0), (1, 1))]
        )
        loss = torch.nn.functional.ctc_loss(
            x, torch.tensor(e), torch.tensor([2]), torch.tensor([1]), reduction="none"
        )
        loss_prob = loss.neg().exp().item()
        path_prob = paths.exp().sum().item()

        torch.testing.assert_allclose(loss_prob, path_prob)
        # Check 1best prob against loss with input_length = 1
        loss = torch.nn.functional.ctc_loss(
            x, torch.tensor(e), torch.tensor([1]), torch.tensor([1]), reduction="none"
        )
        loss_prob = loss.neg().exp().item()
        torch.testing.assert_allclose(
            loss_prob, [p.mean() for p in r["prob-segmentation"]][0].item()
        )
//  [python]
def test_prob(self):
        x = torch.tensor([[[0.3, 0.6, 0.1]], [[0.6, 0.3, 0.2]]]).log()
        decoder = CTCGreedyDecoder()
        r = decoder(x, segmentation=True, apply_softmax=False)
        e = [[1]]
        self.assertEqual(e, r["hyp"])
        # Check actual loss prob
        paths = torch.tensor(
            [x[0, 0, a] + x[1, 0, b] for a, b in ((0, 1), (1, 0), (1, 1))]
        )
        loss = torch.nn.functional.ctc_loss(
            x, torch.tensor(e), torch.tensor([2]), torch.tensor([1]), reduction="none"
        )
        loss_prob = loss.neg().exp().item()
        path_prob = paths.exp().sum().item()

        torch.testing.assert_close(loss_prob, path_prob)
        # Check 1best prob against loss with input_length = 1
        loss = torch.nn.functional.ctc_loss(
            x, torch.tensor(e), torch.tensor([1]), torch.tensor([1]), reduction="none"
        )
        loss_prob = loss.neg().exp().item()
        torch.testing.assert_close(
            loss_prob, [p.mean() for p in r["prob-segmentation"]][0].item()
        )
//  [python]
def test_time_series_loss():
    targets = torch.Tensor(
        np.array(
            [
                [
                    [np.nan, 0.2, np.nan, 0.3, np.nan, 0.1, np.nan, 0.9],
                    [np.nan, 0.8, np.nan, 0.2, np.nan, 0.9, np.nan, 0.2],
                ],
                [
                    [np.nan, 0.3, np.nan, 0.2, np.nan, 0.5, np.nan, 0.1],
                    [np.nan, 0.1, np.nan, 0.2, np.nan, 0.3, np.nan, 0.2],
                ],
            ]
        )
    )
    targets_expected = torch.Tensor(
        np.array([[[0.1, 0.9], [0.9, 0.2]], [[0.5, 0.1], [0.3, 0.2]]])
    )

    preds = torch.Tensor(
        np.array(
            [
                [[0.4, 0.9, 0.4], [0.1, 0.3, 0.8]],
                [[0.2, 0.5, 0.3], [0.2, 0.1, 0.5]],
            ]
        )
    )
    preds_expected = torch.Tensor(
        np.array([[[0.4, 0.4], [0.1, 0.8]], [[0.2, 0.3], [0.2, 0.5]]])
    )

    time_series_loss = TimeSeriesLoss(lambda *args: args)

    preds_out, targets_out = time_series_loss(preds, targets)

    torch.testing.assert_allclose(preds_out, preds_expected.flatten())
    torch.testing.assert_allclose(targets_out, targets_expected.flatten())
//  [python]
def test_time_series_loss():
    targets = torch.Tensor(
        np.array(
            [
                [
                    [np.nan, 0.2, np.nan, 0.3, np.nan, 0.1, np.nan, 0.9],
                    [np.nan, 0.8, np.nan, 0.2, np.nan, 0.9, np.nan, 0.2],
                ],
                [
                    [np.nan, 0.3, np.nan, 0.2, np.nan, 0.5, np.nan, 0.1],
                    [np.nan, 0.1, np.nan, 0.2, np.nan, 0.3, np.nan, 0.2],
                ],
            ]
        )
    )
    targets_expected = torch.Tensor(
        np.array([[[0.1, 0.9], [0.9, 0.2]], [[0.5, 0.1], [0.3, 0.2]]])
    )

    preds = torch.Tensor(
        np.array(
            [
                [[0.4, 0.9, 0.4], [0.1, 0.3, 0.8]],
                [[0.2, 0.5, 0.3], [0.2, 0.1, 0.5]],
            ]
        )
    )
    preds_expected = torch.Tensor(
        np.array([[[0.4, 0.4], [0.1, 0.8]], [[0.2, 0.3], [0.2, 0.5]]])
    )

    time_series_loss = TimeSeriesLoss(lambda *args: args)

    preds_out, targets_out = time_series_loss(preds, targets)

    torch.testing.assert_close(preds_out, preds_expected.flatten())
    torch.testing.assert_close(targets_out, targets_expected.flatten())
//  [python]
def test_scatter_topk(device):
    src = torch.tensor([
        2.3, 8.3, 1.4,
        3.7, 6.6, 1.7, 8.6, 6.4, 7.6,
        6.8, 9.1,
        5.2, 2.0, 0.4, 1.0
    ], device=device)
    index = torch.tensor([
        0, 0, 0,
        1, 1, 1, 1, 1, 1,
        2, 2,
        3, 3, 3, 3
    ], device=device)

    result_values, result_indexes_whole, result_indexes_within = scatter_topk(src, index, k=4)

    expected_values = torch.tensor([
        8.3, 2.3, 1.4, float('NaN'),
        8.6, 7.6, 6.6, 6.4,
        9.1, 6.8, float('NaN'), float('NaN'),
        5.2, 2.0, 1.0, 0.4,
    ], device=device)
    expected_indexes_whole = torch.tensor([
         1,  0,  2, -1,
         6,  8,  4,  7,
        10,  9, -1, -1,
        11, 12, 14, 13,
    ], device=device)
    expected_indexes_within = torch.tensor([
        1,  0,  2, -1,
        3,  5,  1,  4,
        1,  0, -1, -1,
        0,  1,  3,  2,
    ], device=device)

    torch.testing.assert_allclose(result_values, expected_values, equal_nan=True)
    torch.testing.assert_allclose(result_indexes_whole, expected_indexes_whole)
    torch.testing.assert_allclose(result_indexes_within, expected_indexes_within)
//  [python]
def test_scatter_topk(device):
    src = torch.tensor([
        2.3, 8.3, 1.4,
        3.7, 6.6, 1.7, 8.6, 6.4, 7.6,
        6.8, 9.1,
        5.2, 2.0, 0.4, 1.0
    ], device=device)
    index = torch.tensor([
        0, 0, 0,
        1, 1, 1, 1, 1, 1,
        2, 2,
        3, 3, 3, 3
    ], device=device)

    result_values, result_indexes_whole, result_indexes_within = scatter_topk(src, index, k=4)

    expected_values = torch.tensor([
        8.3, 2.3, 1.4, float('NaN'),
        8.6, 7.6, 6.6, 6.4,
        9.1, 6.8, float('NaN'), float('NaN'),
        5.2, 2.0, 1.0, 0.4,
    ], device=device)
    expected_indexes_whole = torch.tensor([
         1,  0,  2, -1,
         6,  8,  4,  7,
        10,  9, -1, -1,
        11, 12, 14, 13,
    ], device=device)
    expected_indexes_within = torch.tensor([
        1,  0,  2, -1,
        3,  5,  1,  4,
        1,  0, -1, -1,
        0,  1,  3,  2,
    ], device=device)

    torch.testing.assert_close(result_values, expected_values, equal_nan=True)
    torch.testing.assert_close(result_indexes_whole, expected_indexes_whole)
    torch.testing.assert_close(result_indexes_within, expected_indexes_within)
//  [python]
def _test_non_randomness(self, device):
        # Noises should be the same in a batch
        x0 = torch.normal(0, 1, size=(1, 6), dtype=torch.float32, device=device)
        x = x0.repeat(2, 1)
        y = self.linear(x)
        torch.testing.assert_allclose(y[0], y[1], rtol=1e-4, atol=0)
//  [python]
def _test_non_randomness(self, device):
        # Noises should be the same in a batch
        x0 = torch.normal(0, 1, size=(1, 6), dtype=torch.float32, device=device)
        x = x0.repeat(2, 1)
        y = self.linear(x)
        torch.testing.assert_close(y[0], y[1], rtol=1e-4, atol=0)
//  [python]
def test_poisson_log_prob(self):
        rate = torch.randn(2, 3).abs().requires_grad_()
        rate_1d = torch.randn(1).abs().requires_grad_()
        rate_zero = torch.zeros([], requires_grad=True)

        def ref_log_prob(ref_rate, idx, x, log_prob):
            l = ref_rate.view(-1)[idx].detach()
            expected = scipy.stats.poisson.logpmf(x, l)
            self.assertEqual(log_prob, expected, atol=1e-3, rtol=0)

        set_rng_seed(0)
        self._check_log_prob(Poisson(rate), lambda *args: ref_log_prob(rate, *args))
        self._check_log_prob(Poisson(rate_zero), lambda *args: ref_log_prob(rate_zero, *args))
        self._gradcheck_log_prob(Poisson, (rate,))
        self._gradcheck_log_prob(Poisson, (rate_1d,))

        # We cannot check gradients automatically for zero rates because the finite difference
        # approximation enters the forbidden parameter space. We instead compare with the
        # theoretical results.
        dist = Poisson(rate_zero)
        s = dist.sample()
        dist.log_prob(s).backward()
        torch.testing.assert_allclose(rate_zero.grad, -1.0)
        dist.log_prob(torch.ones_like(rate_zero)).backward()
        torch.testing.assert_allclose(rate_zero.grad, torch.inf)
//  [python]
def test_poisson_log_prob(self):
        rate = torch.randn(2, 3).abs().requires_grad_()
        rate_1d = torch.randn(1).abs().requires_grad_()
        rate_zero = torch.zeros([], requires_grad=True)

        def ref_log_prob(ref_rate, idx, x, log_prob):
            l = ref_rate.view(-1)[idx].detach()
            expected = scipy.stats.poisson.logpmf(x, l)
            self.assertEqual(log_prob, expected, atol=1e-3, rtol=0)

        set_rng_seed(0)
        self._check_log_prob(Poisson(rate), lambda *args: ref_log_prob(rate, *args))
        self._check_log_prob(Poisson(rate_zero), lambda *args: ref_log_prob(rate_zero, *args))
        self._gradcheck_log_prob(Poisson, (rate,))
        self._gradcheck_log_prob(Poisson, (rate_1d,))

        # We cannot check gradients automatically for zero rates because the finite difference
        # approximation enters the forbidden parameter space. We instead compare with the
        # theoretical results.
        dist = Poisson(rate_zero)
        s = dist.sample()
        dist.log_prob(s).backward()
        torch.testing.assert_close(rate_zero.grad, -1.0)
        dist.log_prob(torch.ones_like(rate_zero)).backward()
        torch.testing.assert_close(rate_zero.grad, torch.inf)
//  [python]
def _load_local_and_train(self, config, rank, group, d_model=16, d_vocab=23):
        """Check that local_state_dict can be saved and loaded for a given worker, and that training updates it"""
        model = self.get_wrapped_model(
            group, cuda_first=False, config=config, d_vocab=d_vocab, d_model=d_model, add_bn=False
        )  # Set bn=True here to show that BN doesn't get updated
        state_1 = model.local_state_dict()
        state_before_training = {k: v.cpu().clone() for k, v in state_1.items()}
        assert len(state_1) > 0
        model.load_local_state_dict(state_1)
        weight_key = "flat_param_0" if model.flatten_parameters else "embed_tokens.weight"

        state_1_weight = state_1[weight_key]
        assert state_1_weight.dtype == torch.float32, f"got dtype {state_1_weight.dtype} expected torch.float32"
        if not model.flatten_parameters:
            # The weight will be sharded since we access module.state_dict directly
            state_1_module_weight = model.module.state_dict()[weight_key]
            torch.testing.assert_allclose(state_1_weight, state_1_module_weight)
            torch.testing.assert_allclose(state_1_weight, model.module.embed_tokens.weight)
        # increasing number of epochs from 1 to 6 for ShardedGradScaler to work properly.
        # test fails for num_epochs < 6 since the updates are skipped due to gradient being inf.
        self._train_for_several_steps(model, 6, model.mixed_precision)

        state_2 = model.local_state_dict()
        state_after_training = {k: v.cpu().clone() for k, v in state_2.items()}
        model.load_local_state_dict(state_2)

        assert state_1.keys() == state_2.keys()

        # Assert that parameters were updated since before training
        unchanged = []
        unwrapped_model = model.module.module
        buffers = {name for name, _ in unwrapped_model.named_buffers()}
        for k in state_1:
            if (state_before_training[k] == state_after_training[k]).all() and (k not in buffers):
                unchanged.append(k)
        if unchanged:
            raise AssertionError(f"params {unchanged} not changed after training")
//  [python]
def _load_local_and_train(self, config, rank, group, d_model=16, d_vocab=23):
        """Check that local_state_dict can be saved and loaded for a given worker, and that training updates it"""
        model = self.get_wrapped_model(
            group, cuda_first=False, config=config, d_vocab=d_vocab, d_model=d_model, add_bn=False
        )  # Set bn=True here to show that BN doesn't get updated
        state_1 = model.local_state_dict()
        state_before_training = {k: v.cpu().clone() for k, v in state_1.items()}
        assert len(state_1) > 0
        model.load_local_state_dict(state_1)
        weight_key = "flat_param_0" if model.flatten_parameters else "embed_tokens.weight"

        state_1_weight = state_1[weight_key]
        assert state_1_weight.dtype == torch.float32, f"got dtype {state_1_weight.dtype} expected torch.float32"
        if not model.flatten_parameters:
            # The weight will be sharded since we access module.state_dict directly
            state_1_module_weight = model.module.state_dict()[weight_key]
            torch.testing.assert_close(state_1_weight, state_1_module_weight)
            torch.testing.assert_close(state_1_weight, model.module.embed_tokens.weight)
        # increasing number of epochs from 1 to 6 for ShardedGradScaler to work properly.
        # test fails for num_epochs < 6 since the updates are skipped due to gradient being inf.
        self._train_for_several_steps(model, 6, model.mixed_precision)

        state_2 = model.local_state_dict()
        state_after_training = {k: v.cpu().clone() for k, v in state_2.items()}
        model.load_local_state_dict(state_2)

        assert state_1.keys() == state_2.keys()

        # Assert that parameters were updated since before training
        unchanged = []
        unwrapped_model = model.module.module
        buffers = {name for name, _ in unwrapped_model.named_buffers()}
        for k in state_1:
            if (state_before_training[k] == state_after_training[k]).all() and (k not in buffers):
                unchanged.append(k)
        if unchanged:
            raise AssertionError(f"params {unchanged} not changed after training")
//  [python]
def test_vqvae(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        set_seed(0)
        x = torch.rand((1, 5120, 1))
        with torch.no_grad():
            zs = model.vqvae.encode(x, start_level=2, bs_chunks=x.shape[0])
        torch.testing.assert_allclose(zs[0][0], torch.tensor(self.EXPECTED_VQVAE_ENCODE))

        with torch.no_grad():
            x = model.vqvae.decode(zs, start_level=2, bs_chunks=x.shape[0])
        torch.testing.assert_allclose(x[0, :40, 0], torch.tensor(self.EXPECTED_VQVAE_DECODE), atol=1e-4, rtol=1e-4)
//  [python]
def test_vqvae(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        set_seed(0)
        x = torch.rand((1, 5120, 1))
        with torch.no_grad():
            zs = model.vqvae.encode(x, start_level=2, bs_chunks=x.shape[0])
        torch.testing.assert_close(zs[0][0], torch.tensor(self.EXPECTED_VQVAE_ENCODE))

        with torch.no_grad():
            x = model.vqvae.decode(zs, start_level=2, bs_chunks=x.shape[0])
        torch.testing.assert_close(x[0, :40, 0], torch.tensor(self.EXPECTED_VQVAE_DECODE), atol=1e-4, rtol=1e-4)
//  [python]
def test_multiply():
    n = 10
    x_np = np.random.rand(n) + 1j * np.random.rand(n)
    y_np = np.random.rand(n) + 1j * np.random.rand(n)
    xy_np = x_np * y_np
    xystack_np = np.stack([np.real(xy_np), np.imag(xy_np)], axis=-1)
    x = torch.stack([torch.from_numpy(np.real(x_np)), torch.from_numpy(np.imag(x_np))], dim=-1)
    y = torch.stack([torch.from_numpy(np.real(y_np)), torch.from_numpy(np.imag(y_np))], dim=-1)
    xystack = complex.multiply(x, y)
    torch.testing.assert_allclose(xystack_np, xystack)
//  [python]
def test_multiply():
    n = 10
    x_np = np.random.rand(n) + 1j * np.random.rand(n)
    y_np = np.random.rand(n) + 1j * np.random.rand(n)
    xy_np = x_np * y_np
    xystack_np = np.stack([np.real(xy_np), np.imag(xy_np)], axis=-1)
    x = torch.stack([torch.from_numpy(np.real(x_np)), torch.from_numpy(np.imag(x_np))], dim=-1)
    y = torch.stack([torch.from_numpy(np.real(y_np)), torch.from_numpy(np.imag(y_np))], dim=-1)
    xystack = complex.multiply(x, y)
    torch.testing.assert_close(xystack_np, xystack)
//  [python]
def test_conditioning(self):
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()

        labels = self.prepare_inputs()
        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long) for _ in range(3)]

        top_prior = model.priors[0]
        start = 0
        music_token_conds = top_prior.get_music_tokens_conds(zs, start=start, end=start + top_prior.n_ctx)
        metadata = top_prior.get_metadata(labels[0].clone(), start, 1058304, 0)

        self.assertIsNone(music_token_conds)
        self.assertListEqual(metadata.numpy()[0][:10].tolist(), self.EXPECTED_Y_COND)

        audio_conditioning, metadata_conditioning, lyric_tokens = top_prior.get_cond(music_token_conds, metadata)
        torch.testing.assert_allclose(
            audio_conditioning[0][0][:30].detach(), torch.tensor(self.EXPECTED_AUDIO_COND), atol=1e-4, rtol=1e-4
        )
        torch.testing.assert_allclose(
            metadata_conditioning[0][0][:30].detach(), torch.tensor(self.EXPECTED_META_COND), atol=1e-4, rtol=1e-4
        )
        torch.testing.assert_allclose(
            lyric_tokens[0, :30].detach(), torch.tensor(self.EXPECTED_LYRIC_COND), atol=1e-4, rtol=1e-4
        )
//  [python]
def test_conditioning(self):
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()

        labels = self.prepare_inputs()
        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long) for _ in range(3)]

        top_prior = model.priors[0]
        start = 0
        music_token_conds = top_prior.get_music_tokens_conds(zs, start=start, end=start + top_prior.n_ctx)
        metadata = top_prior.get_metadata(labels[0].clone(), start, 1058304, 0)

        self.assertIsNone(music_token_conds)
        self.assertListEqual(metadata.numpy()[0][:10].tolist(), self.EXPECTED_Y_COND)

        audio_conditioning, metadata_conditioning, lyric_tokens = top_prior.get_cond(music_token_conds, metadata)
        torch.testing.assert_close(
            audio_conditioning[0][0][:30].detach(), torch.tensor(self.EXPECTED_AUDIO_COND), atol=1e-4, rtol=1e-4
        )
        torch.testing.assert_close(
            metadata_conditioning[0][0][:30].detach(), torch.tensor(self.EXPECTED_META_COND), atol=1e-4, rtol=1e-4
        )
        torch.testing.assert_close(
            lyric_tokens[0, :30].detach(), torch.tensor(self.EXPECTED_LYRIC_COND), atol=1e-4, rtol=1e-4
        )
//  [python]
def test_conj():
    n = 10
    x_np = np.random.rand(n) + 1j * np.random.rand(n)
    xconj_np = np.conj(x_np)
    x = torch.stack([torch.from_numpy(np.real(x_np)), torch.from_numpy(np.imag(x_np))], dim=-1)
    xconjstack_np = torch.stack([torch.from_numpy(np.real(xconj_np)), torch.from_numpy(np.imag(xconj_np))], dim=-1)
    xconj = complex.conj(x)
    torch.testing.assert_allclose(xconjstack_np, xconj)
//  [python]
def test_conj():
    n = 10
    x_np = np.random.rand(n) + 1j * np.random.rand(n)
    xconj_np = np.conj(x_np)
    x = torch.stack([torch.from_numpy(np.real(x_np)), torch.from_numpy(np.imag(x_np))], dim=-1)
    xconjstack_np = torch.stack([torch.from_numpy(np.real(xconj_np)), torch.from_numpy(np.imag(xconj_np))], dim=-1)
    xconj = complex.conj(x)
    torch.testing.assert_close(xconjstack_np, xconj)
//  [python]
def test_abs2():
    n = 10
    x_np = np.random.rand(n) + 1j * np.random.rand(n)
    xabs2_np = np.abs(x_np) ** 2
    xabs2_np = torch.from_numpy(xabs2_np)
    x = torch.stack([torch.from_numpy(np.real(x_np)), torch.from_numpy(np.imag(x_np))], dim=-1)
    xabs2 = complex.abs2(x)
    torch.testing.assert_allclose(xabs2_np, xabs2)
//  [python]
def test_abs2():
    n = 10
    x_np = np.random.rand(n) + 1j * np.random.rand(n)
    xabs2_np = np.abs(x_np) ** 2
    xabs2_np = torch.from_numpy(xabs2_np)
    x = torch.stack([torch.from_numpy(np.real(x_np)), torch.from_numpy(np.imag(x_np))], dim=-1)
    xabs2 = complex.abs2(x)
    torch.testing.assert_close(xabs2_np, xabs2)
//  [python]
def test_slow_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        labels = [i.to(torch_device) for i in self.prepare_inputs(self.model_id)]

        set_seed(0)
        model.priors[0].to(torch_device)
        zs = [torch.zeros(1, 0, dtype=torch.long).to(torch_device) for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[0][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_2))
        model.priors[0].cpu()

        set_seed(0)
        model.priors[1].to(torch_device)
        zs = model._sample(zs, labels, [1], sample_length=60 * model.priors[1].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[1][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_1))
        model.priors[1].cpu()

        set_seed(0)
        model.priors[2].to(torch_device)
        zs = model._sample(zs, labels, [2], sample_length=60 * model.priors[2].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[2][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_0))
//  [python]
def test_slow_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        labels = [i.to(torch_device) for i in self.prepare_inputs(self.model_id)]

        set_seed(0)
        model.priors[0].to(torch_device)
        zs = [torch.zeros(1, 0, dtype=torch.long).to(torch_device) for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[0][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_2))
        model.priors[0].cpu()

        set_seed(0)
        model.priors[1].to(torch_device)
        zs = model._sample(zs, labels, [1], sample_length=60 * model.priors[1].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[1][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_1))
        model.priors[1].cpu()

        set_seed(0)
        model.priors[2].to(torch_device)
        zs = model._sample(zs, labels, [2], sample_length=60 * model.priors[2].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[2][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_0))
//  [python]
def test_primed_sampling(self):
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False

        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        set_seed(0)
        waveform = torch.rand((1, 5120, 1))
        tokens = list(self.prepare_inputs())

        zs = [model.vqvae.encode(waveform, start_level=2, bs_chunks=waveform.shape[0])[0], None, None]
        zs = model._sample(
            zs, tokens, sample_levels=[0], save_results=False, sample_length=40 * model.priors[0].raw_to_tokens
        )
        torch.testing.assert_allclose(zs[0][0][:40], torch.tensor(self.EXPECTED_PRIMED_0))

        upper_2 = torch.cat((zs[0], torch.zeros(1, 2048 - zs[0].shape[-1])), dim=-1).long()
        zs = [upper_2, model.vqvae.encode(waveform, start_level=1, bs_chunks=waveform.shape[0])[0], None]
        zs = model._sample(
            zs, tokens, sample_levels=[1], save_results=False, sample_length=40 * model.priors[1].raw_to_tokens
        )
        torch.testing.assert_allclose(zs[1][0][:40], torch.tensor(self.EXPECTED_PRIMED_1))

        upper_1 = torch.cat((zs[1], torch.zeros(1, 2048 - zs[1].shape[-1])), dim=-1).long()
        zs = [upper_2, upper_1, model.vqvae.encode(waveform, start_level=0, bs_chunks=waveform.shape[0])[0]]
        zs = model._sample(
            zs, tokens, sample_levels=[2], save_results=False, sample_length=40 * model.priors[2].raw_to_tokens
        )
        torch.testing.assert_allclose(zs[2][0][:40].cpu(), torch.tensor(self.EXPECTED_PRIMED_2))
//  [python]
def test_primed_sampling(self):
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False

        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        set_seed(0)
        waveform = torch.rand((1, 5120, 1))
        tokens = list(self.prepare_inputs())

        zs = [model.vqvae.encode(waveform, start_level=2, bs_chunks=waveform.shape[0])[0], None, None]
        zs = model._sample(
            zs, tokens, sample_levels=[0], save_results=False, sample_length=40 * model.priors[0].raw_to_tokens
        )
        torch.testing.assert_close(zs[0][0][:40], torch.tensor(self.EXPECTED_PRIMED_0))

        upper_2 = torch.cat((zs[0], torch.zeros(1, 2048 - zs[0].shape[-1])), dim=-1).long()
        zs = [upper_2, model.vqvae.encode(waveform, start_level=1, bs_chunks=waveform.shape[0])[0], None]
        zs = model._sample(
            zs, tokens, sample_levels=[1], save_results=False, sample_length=40 * model.priors[1].raw_to_tokens
        )
        torch.testing.assert_close(zs[1][0][:40], torch.tensor(self.EXPECTED_PRIMED_1))

        upper_1 = torch.cat((zs[1], torch.zeros(1, 2048 - zs[1].shape[-1])), dim=-1).long()
        zs = [upper_2, upper_1, model.vqvae.encode(waveform, start_level=0, bs_chunks=waveform.shape[0])[0]]
        zs = model._sample(
            zs, tokens, sample_levels=[2], save_results=False, sample_length=40 * model.priors[2].raw_to_tokens
        )
        torch.testing.assert_close(zs[2][0][:40].cpu(), torch.tensor(self.EXPECTED_PRIMED_2))
//  [python]
def test_t5_layer_norm(batch_size=16, hidden_size=128, seq_len=4):
    t5_layer_norm = transformers.models.t5.modeling_t5.T5LayerNorm(hidden_size=hidden_size).to(DEVICE)
    l1, l2 = tuple(copy.deepcopy(t5_layer_norm) for _ in range(2))

    inputs = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE)
    targets = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE)

    grad1 = []
    for i, t in utils.zip_(inputs, targets):
        i, t = i[None, :], t[None, :]
        l1.zero_grad()
        (.5 * (l1(i) - t) ** 2.).sum().backward()
        grad1.append(l1.weight.grad.detach().clone())
    grad1 = torch.stack(grad1)

    l2.zero_grad()
    outputs = l2(inputs)
    grad_outputs = (outputs - targets)
    supported_layers_grad_samplers._compute_t5_layer_norm_grad_sample(l2, (inputs,), (grad_outputs,))
    grad2 = l2.weight.grad_sample
    torch.testing.assert_allclose(grad1, grad2, atol=1e-5, rtol=1e-4)
//  [python]
def test_t5_layer_norm(batch_size=16, hidden_size=128, seq_len=4):
    t5_layer_norm = transformers.models.t5.modeling_t5.T5LayerNorm(hidden_size=hidden_size).to(DEVICE)
    l1, l2 = tuple(copy.deepcopy(t5_layer_norm) for _ in range(2))

    inputs = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE)
    targets = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE)

    grad1 = []
    for i, t in utils.zip_(inputs, targets):
        i, t = i[None, :], t[None, :]
        l1.zero_grad()
        (.5 * (l1(i) - t) ** 2.).sum().backward()
        grad1.append(l1.weight.grad.detach().clone())
    grad1 = torch.stack(grad1)

    l2.zero_grad()
    outputs = l2(inputs)
    grad_outputs = (outputs - targets)
    supported_layers_grad_samplers._compute_t5_layer_norm_grad_sample(l2, (inputs,), (grad_outputs,))
    grad2 = l2.weight.grad_sample
    torch.testing.assert_close(grad1, grad2, atol=1e-5, rtol=1e-4)
//  [python]
def assertTensorEqual(self,
                          tensor: torch.Tensor,
                          value: float,
                          msg=None,
                          **kwarg):
        tensor = tensor.to(torch.float32)
        value = torch.FloatTensor([value])
        try:
            torch.testing.assert_allclose(tensor, value, **kwarg)
        except AssertionError as e:
            self.fail(self._formatMessage(msg, str(e)))
//  [python]
def assertTensorEqual(self,
                          tensor: torch.Tensor,
                          value: float,
                          msg=None,
                          **kwarg):
        tensor = tensor.to(torch.float32)
        value = torch.FloatTensor([value])
        try:
            torch.testing.assert_close(tensor, value, **kwarg)
        except AssertionError as e:
            self.fail(self._formatMessage(msg, str(e)))
//  [python]
def assertTensorEqual(self,
                          tensor: torch.Tensor,
                          value: float,
                          msg=None,
                          **kwarg):
        tensor = tensor.to(torch.float32)
        value = torch.tensor(value).float()
        try:
            torch.testing.assert_allclose(tensor, value, **kwarg)
        except AssertionError as e:
            self.fail(self._formatMessage(msg, str(e)))
//  [python]
def assertTensorEqual(self,
                          tensor: torch.Tensor,
                          value: float,
                          msg=None,
                          **kwarg):
        tensor = tensor.to(torch.float32)
        value = torch.tensor(value).float()
        try:
            torch.testing.assert_close(tensor, value, **kwarg)
        except AssertionError as e:
            self.fail(self._formatMessage(msg, str(e)))
//  [python]
def test_conditioning(self):
        torch.backends.cuda.matmul.allow_tf32 = False
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()

        labels = self.prepare_inputs()
        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long) for _ in range(3)]

        top_prior = model.priors[0]
        start = 0
        music_token_conds = top_prior.get_music_tokens_conds(zs, start=start, end=start + top_prior.n_ctx)
        metadata = top_prior.get_metadata(labels[0].clone(), start, 1058304, 0)

        self.assertIsNone(music_token_conds)
        self.assertListEqual(metadata.numpy()[0][:10].tolist(), self.EXPECTED_Y_COND)

        audio_conditioning, metadata_conditioning, lyric_tokens = top_prior.get_cond(music_token_conds, metadata)
        torch.testing.assert_allclose(
            audio_conditioning[0][0][:30].detach(), torch.tensor(self.EXPECTED_AUDIO_COND), atol=1e-4, rtol=1e-4
        )
        torch.testing.assert_allclose(
            metadata_conditioning[0][0][:30].detach(), torch.tensor(self.EXPECTED_META_COND), atol=1e-4, rtol=1e-4
        )
        torch.testing.assert_allclose(
            lyric_tokens[0, :30].detach(), torch.tensor(self.EXPECTED_LYRIC_COND), atol=1e-4, rtol=1e-4
        )
//  [python]
def test_conditioning(self):
        torch.backends.cuda.matmul.allow_tf32 = False
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()

        labels = self.prepare_inputs()
        set_seed(0)
        zs = [torch.zeros(1, 0, dtype=torch.long) for _ in range(3)]

        top_prior = model.priors[0]
        start = 0
        music_token_conds = top_prior.get_music_tokens_conds(zs, start=start, end=start + top_prior.n_ctx)
        metadata = top_prior.get_metadata(labels[0].clone(), start, 1058304, 0)

        self.assertIsNone(music_token_conds)
        self.assertListEqual(metadata.numpy()[0][:10].tolist(), self.EXPECTED_Y_COND)

        audio_conditioning, metadata_conditioning, lyric_tokens = top_prior.get_cond(music_token_conds, metadata)
        torch.testing.assert_close(
            audio_conditioning[0][0][:30].detach(), torch.tensor(self.EXPECTED_AUDIO_COND), atol=1e-4, rtol=1e-4
        )
        torch.testing.assert_close(
            metadata_conditioning[0][0][:30].detach(), torch.tensor(self.EXPECTED_META_COND), atol=1e-4, rtol=1e-4
        )
        torch.testing.assert_close(
            lyric_tokens[0, :30].detach(), torch.tensor(self.EXPECTED_LYRIC_COND), atol=1e-4, rtol=1e-4
        )
//  [python]
def compare_tensors(self, tensors1, tensors2, rtol=1e-3, atol=1e-7):
        for tensor1, tensor2 in zip(tensors1, tensors2):
            torch.testing.assert_allclose(tensor1, tensor2, rtol=rtol, atol=atol)
//  [python]
def compare_tensors(self, tensors1, tensors2, rtol=1e-3, atol=1e-7):
        for tensor1, tensor2 in zip(tensors1, tensors2):
            torch.testing.assert_close(tensor1, tensor2, rtol=rtol, atol=atol)
//  [python]
def test_unique_indices(
    indices: Tensor, indices_: Tensor, inv_: Tensor
) -> None:
    """Unit test for computing unique indices of a tensor."""
    inv = unique_indices(indices)
    torch.testing.assert_allclose(indices[inv], indices_)
    torch.testing.assert_allclose(inv, inv_)
//  [python]
def test_unique_indices(
    indices: Tensor, indices_: Tensor, inv_: Tensor
) -> None:
    """Unit test for computing unique indices of a tensor."""
    inv = unique_indices(indices)
    torch.testing.assert_close(indices[inv], indices_)
    torch.testing.assert_close(inv, inv_)
//  [python]
def test_slow_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        labels = [i.cuda() for i in self.prepare_inputs(self.model_id)]

        set_seed(0)
        model.priors[0].cuda()
        zs = [torch.zeros(1, 0, dtype=torch.long).cuda() for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[0][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_2))
        model.priors[0].cpu()

        set_seed(0)
        model.priors[1].cuda()
        zs = model._sample(zs, labels, [1], sample_length=60 * model.priors[1].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[1][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_1))
        model.priors[1].cpu()

        set_seed(0)
        model.priors[2].cuda()
        zs = model._sample(zs, labels, [2], sample_length=60 * model.priors[2].raw_to_tokens, save_results=False)
        torch.testing.assert_allclose(zs[2][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_0))
//  [python]
def test_slow_sampling(self):
        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        labels = [i.cuda() for i in self.prepare_inputs(self.model_id)]

        set_seed(0)
        model.priors[0].cuda()
        zs = [torch.zeros(1, 0, dtype=torch.long).cuda() for _ in range(3)]
        zs = model._sample(zs, labels, [0], sample_length=60 * model.priors[0].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[0][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_2))
        model.priors[0].cpu()

        set_seed(0)
        model.priors[1].cuda()
        zs = model._sample(zs, labels, [1], sample_length=60 * model.priors[1].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[1][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_1))
        model.priors[1].cpu()

        set_seed(0)
        model.priors[2].cuda()
        zs = model._sample(zs, labels, [2], sample_length=60 * model.priors[2].raw_to_tokens, save_results=False)
        torch.testing.assert_close(zs[2][0].cpu(), torch.tensor(self.EXPECTED_GPU_OUTPUTS_0))
//  [python]
def test_voxel_pool() -> None:
    """Test voxel pooling method."""
    voxels = _construct_dummy_voxels()

    # Grab indices and values.
    indices = voxels["indices"]
    values = voxels["values"]

    # Declare voxel grid size.
    size = [6, 6, 6]

    # Apply mean pooling.
    indices, values, _ = cluster_grid(
        indices, values, size, cluster_type=ClusterType.MEAN
    )

    # Declare expected indices and values.
    indices_expected = torch.as_tensor(
        [[0, 0, 0], [5, 5, 5]], dtype=torch.int64
    )
    values_expected = torch.as_tensor(
        [[1.1, 1.1, 1.1], [5.35, 5.35, 5.35]], dtype=torch.float32
    )

    # Check if all indices and values match.
    torch.testing.assert_allclose(indices_expected, indices)
    torch.testing.assert_allclose(values_expected, values)
//  [python]
def test_voxel_pool() -> None:
    """Test voxel pooling method."""
    voxels = _construct_dummy_voxels()

    # Grab indices and values.
    indices = voxels["indices"]
    values = voxels["values"]

    # Declare voxel grid size.
    size = [6, 6, 6]

    # Apply mean pooling.
    indices, values, _ = cluster_grid(
        indices, values, size, cluster_type=ClusterType.MEAN
    )

    # Declare expected indices and values.
    indices_expected = torch.as_tensor(
        [[0, 0, 0], [5, 5, 5]], dtype=torch.int64
    )
    values_expected = torch.as_tensor(
        [[1.1, 1.1, 1.1], [5.35, 5.35, 5.35]], dtype=torch.float32
    )

    # Check if all indices and values match.
    torch.testing.assert_close(indices_expected, indices)
    torch.testing.assert_close(values_expected, values)
//  [python]
def test_primed_sampling(self):
        torch.backends.cuda.matmul.allow_tf32 = False

        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        set_seed(0)
        waveform = torch.rand((1, 5120, 1))
        tokens = list(self.prepare_inputs())

        zs = [model.vqvae.encode(waveform, start_level=2, bs_chunks=waveform.shape[0])[0], None, None]
        zs = model._sample(
            zs, tokens, sample_levels=[0], save_results=False, sample_length=40 * model.priors[0].raw_to_tokens
        )
        torch.testing.assert_allclose(zs[0][0][:40], torch.tensor(self.EXPECTED_PRIMED_0))

        upper_2 = torch.cat((zs[0], torch.zeros(1, 2048 - zs[0].shape[-1])), dim=-1).long()
        zs = [upper_2, model.vqvae.encode(waveform, start_level=1, bs_chunks=waveform.shape[0])[0], None]
        zs = model._sample(
            zs, tokens, sample_levels=[1], save_results=False, sample_length=40 * model.priors[1].raw_to_tokens
        )
        torch.testing.assert_allclose(zs[1][0][:40], torch.tensor(self.EXPECTED_PRIMED_1))

        upper_1 = torch.cat((zs[1], torch.zeros(1, 2048 - zs[1].shape[-1])), dim=-1).long()
        zs = [upper_2, upper_1, model.vqvae.encode(waveform, start_level=0, bs_chunks=waveform.shape[0])[0]]
        zs = model._sample(
            zs, tokens, sample_levels=[2], save_results=False, sample_length=40 * model.priors[2].raw_to_tokens
        )
        torch.testing.assert_allclose(zs[2][0][:40].cpu(), torch.tensor(self.EXPECTED_PRIMED_2))
//  [python]
def test_primed_sampling(self):
        torch.backends.cuda.matmul.allow_tf32 = False

        model = JukeboxModel.from_pretrained(self.model_id, min_duration=0).eval()
        set_seed(0)
        waveform = torch.rand((1, 5120, 1))
        tokens = list(self.prepare_inputs())

        zs = [model.vqvae.encode(waveform, start_level=2, bs_chunks=waveform.shape[0])[0], None, None]
        zs = model._sample(
            zs, tokens, sample_levels=[0], save_results=False, sample_length=40 * model.priors[0].raw_to_tokens
        )
        torch.testing.assert_close(zs[0][0][:40], torch.tensor(self.EXPECTED_PRIMED_0))

        upper_2 = torch.cat((zs[0], torch.zeros(1, 2048 - zs[0].shape[-1])), dim=-1).long()
        zs = [upper_2, model.vqvae.encode(waveform, start_level=1, bs_chunks=waveform.shape[0])[0], None]
        zs = model._sample(
            zs, tokens, sample_levels=[1], save_results=False, sample_length=40 * model.priors[1].raw_to_tokens
        )
        torch.testing.assert_close(zs[1][0][:40], torch.tensor(self.EXPECTED_PRIMED_1))

        upper_1 = torch.cat((zs[1], torch.zeros(1, 2048 - zs[1].shape[-1])), dim=-1).long()
        zs = [upper_2, upper_1, model.vqvae.encode(waveform, start_level=0, bs_chunks=waveform.shape[0])[0]]
        zs = model._sample(
            zs, tokens, sample_levels=[2], save_results=False, sample_length=40 * model.priors[2].raw_to_tokens
        )
        torch.testing.assert_close(zs[2][0][:40].cpu(), torch.tensor(self.EXPECTED_PRIMED_2))
//  [python]
def test_kernel_gaussian(self, task_name="gauss"):

        keops_kernel = LazyKeopsKernel(kernel_type="gauss", sigma=0.1)
        torch_kernel = TorchKernel(kernel_type="gauss", sigma=0.1)
        keops_kernel = timming(
            keops_kernel, "test_kernel_{} with keops".format(task_name)
        )
        torch_kernel = timming(
            torch_kernel, "test_kernel_{} with torch".format(task_name)
        )
        keops_gauss = keops_kernel(self.x, self.y, self.b)
        torch_gauss = torch_kernel(self.x, self.y, self.b)
        grad_keops = timming(grad, "test_kernel_{} grad with keops".format(task_name))
        grad_torch = timming(grad, "test_kernel_{} grad with torch".format(task_name))
        keops_grads = grad_keops(
            keops_gauss.mean(), (self.x, self.y, self.b), retain_graph=True
        )
        torch_grads = grad_torch(
            torch_gauss.mean(), (self.x, self.y, self.b), retain_graph=True
        )
        torch.testing.assert_allclose(keops_gauss, torch_gauss, rtol=1e-3, atol=1e-7)
        self.compare_tensors(keops_grads, torch_grads, rtol=1e-3, atol=1e-7)
//  [python]
def test_kernel_gaussian(self, task_name="gauss"):

        keops_kernel = LazyKeopsKernel(kernel_type="gauss", sigma=0.1)
        torch_kernel = TorchKernel(kernel_type="gauss", sigma=0.1)
        keops_kernel = timming(
            keops_kernel, "test_kernel_{} with keops".format(task_name)
        )
        torch_kernel = timming(
            torch_kernel, "test_kernel_{} with torch".format(task_name)
        )
        keops_gauss = keops_kernel(self.x, self.y, self.b)
        torch_gauss = torch_kernel(self.x, self.y, self.b)
        grad_keops = timming(grad, "test_kernel_{} grad with keops".format(task_name))
        grad_torch = timming(grad, "test_kernel_{} grad with torch".format(task_name))
        keops_grads = grad_keops(
            keops_gauss.mean(), (self.x, self.y, self.b), retain_graph=True
        )
        torch_grads = grad_torch(
            torch_gauss.mean(), (self.x, self.y, self.b), retain_graph=True
        )
        torch.testing.assert_close(keops_gauss, torch_gauss, rtol=1e-3, atol=1e-7)
        self.compare_tensors(keops_grads, torch_grads, rtol=1e-3, atol=1e-7)
//  [python]
def test_kernel_gaussian_grad(self, task_name="gauss_grad"):
        keops_kernel = LazyKeopsKernel(kernel_type="gauss_grad", sigma=0.1)
        torch_kernel = TorchKernel(kernel_type="gauss_grad", sigma=0.1)
        keops_kernel = timming(
            keops_kernel, "test_kernel_{} with keops".format(task_name)
        )
        torch_kernel = timming(
            torch_kernel, "test_kernel_{} with torch".format(task_name)
        )
        keops_gauss = keops_kernel(self.px, self.x, self.py, self.y)
        torch_gauss = torch_kernel(self.px, self.x, self.py, self.y)
        grad_keops = timming(grad, "test_kernel_{} grad with keops".format(task_name))
        grad_torch = timming(grad, "test_kernel_{} grad with torch".format(task_name))
        keops_grads = grad_keops(
            keops_gauss.mean(), (self.px, self.x, self.py, self.y), retain_graph=True
        )
        torch_grads = grad_torch(
            torch_gauss.mean(), (self.px, self.x, self.py, self.y), retain_graph=True
        )
        torch.testing.assert_allclose(keops_gauss, torch_gauss, rtol=1e-2, atol=1e-5)
        self.compare_tensors(keops_grads, torch_grads, rtol=1e-3, atol=1e-7)
//  [python]
def test_kernel_gaussian_grad(self, task_name="gauss_grad"):
        keops_kernel = LazyKeopsKernel(kernel_type="gauss_grad", sigma=0.1)
        torch_kernel = TorchKernel(kernel_type="gauss_grad", sigma=0.1)
        keops_kernel = timming(
            keops_kernel, "test_kernel_{} with keops".format(task_name)
        )
        torch_kernel = timming(
            torch_kernel, "test_kernel_{} with torch".format(task_name)
        )
        keops_gauss = keops_kernel(self.px, self.x, self.py, self.y)
        torch_gauss = torch_kernel(self.px, self.x, self.py, self.y)
        grad_keops = timming(grad, "test_kernel_{} grad with keops".format(task_name))
        grad_torch = timming(grad, "test_kernel_{} grad with torch".format(task_name))
        keops_grads = grad_keops(
            keops_gauss.mean(), (self.px, self.x, self.py, self.y), retain_graph=True
        )
        torch_grads = grad_torch(
            torch_gauss.mean(), (self.px, self.x, self.py, self.y), retain_graph=True
        )
        torch.testing.assert_close(keops_gauss, torch_gauss, rtol=1e-2, atol=1e-5)
        self.compare_tensors(keops_grads, torch_grads, rtol=1e-3, atol=1e-7)
//  [python]
def test_multihead_attention_layer(self):
        HID_DIM = 256
        QUERY_LEN = 8
        BATCH_SIZE = 128
        LAYERS = 3
        HEADS = 8
        DROPOUT = 0.1
        device = torch.device("cpu")
        attention = MultiHeadAttentionLayer(HID_DIM, HEADS, DROPOUT, device).to(device)
        src = torch.randn(BATCH_SIZE, QUERY_LEN, HID_DIM).to(device)
        src_mask = (src > 0)[:, :, 0].unsqueeze(1).unsqueeze(2).to(device)

        attention.eval()
        attention = torch.jit.script(attention)
        attention.eval()
        o_ref = attention(src, src, src, src_mask)

        attention_a = StaticRuntime(attention)
        o_test = attention_a(src, src, src, src_mask)
        for a, b in zip(o_ref, o_test):
            torch.testing.assert_allclose(a, b)
//  [python]
def test_multihead_attention_layer(self):
        HID_DIM = 256
        QUERY_LEN = 8
        BATCH_SIZE = 128
        LAYERS = 3
        HEADS = 8
        DROPOUT = 0.1
        device = torch.device("cpu")
        attention = MultiHeadAttentionLayer(HID_DIM, HEADS, DROPOUT, device).to(device)
        src = torch.randn(BATCH_SIZE, QUERY_LEN, HID_DIM).to(device)
        src_mask = (src > 0)[:, :, 0].unsqueeze(1).unsqueeze(2).to(device)

        attention.eval()
        attention = torch.jit.script(attention)
        attention.eval()
        o_ref = attention(src, src, src, src_mask)

        attention_a = StaticRuntime(attention)
        o_test = attention_a(src, src, src, src_mask)
        for a, b in zip(o_ref, o_test):
            torch.testing.assert_close(a, b)
//  [python]
def test_kernel_mutli_gaussian(self, task_name="multi_gauss"):
        keops_kernel = LazyKeopsKernel(
            kernel_type="multi_gauss",
            sigma_list=[0.01, 0.05, 0.1],
            weight_list=[0.2, 0.3, 0.5],
        )
        torch_kernel = TorchKernel(
            kernel_type="multi_gauss",
            sigma_list=[0.01, 0.05, 0.1],
            weight_list=[0.2, 0.3, 0.5],
        )
        keops_kernel = timming(
            keops_kernel, "test_kernel_{} with keops".format(task_name)
        )
        torch_kernel = timming(
            torch_kernel, "test_kernel_{} with torch".format(task_name)
        )
        keops_gauss = keops_kernel(self.x, self.y, self.b)
        torch_gauss = torch_kernel(self.x, self.y, self.b)
        grad_keops = timming(grad, "test_kernel_{} grad with keops".format(task_name))
        grad_torch = timming(grad, "test_kernel_{} grad with torch".format(task_name))
        keops_grads = grad_keops(
            keops_gauss.mean(), (self.x, self.y, self.b), retain_graph=True
        )
        torch_grads = grad_torch(
            torch_gauss.mean(), (self.x, self.y, self.b), retain_graph=True
        )
        torch.testing.assert_allclose(keops_gauss, torch_gauss, rtol=1e-3, atol=1e-7)
        self.compare_tensors(keops_grads, torch_grads, rtol=1e-3, atol=1e-7)