"""simple docstring"""






import unittest
from unittest import skip

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow
from transformers.trainer_utils import set_seed


if is_torch_available():
       import torch

       from transformers import JukeboxModel, JukeboxPrior, JukeboxTokenizer


@require_torch
class 					UpperCAmelCase		(unittest.TestCase						):



    """simple docstring"""

    _UpperCAmelCase			:Optional[Any]							=				(JukeboxModel,) if is_torch_available() else ()
    _UpperCAmelCase			:int							=				"openai/jukebox-1b-lyrics"
    _UpperCAmelCase			:int							=				{
        "artist": "Zac Brown Band",
        "genres": "Country",
        "lyrics": "I met a traveller from an antique land,\n    Who said \"Two vast and trunkless legs of stone\n    Stand in the desert. . . . Near them, on the sand,\n    Half sunk a shattered visage lies, whose frown,\n    And wrinkled lip, and sneer of cold command,\n    Tell that its sculptor well those passions read\n    Which yet survive, stamped on these lifeless things,\n    The hand that mocked them, and the heart that fed;\n    And on the pedestal, these words appear:\n    My name is Ozymandias, King of Kings;\n    Look on my Works, ye Mighty, and despair!\n    Nothing beside remains. Round the decay\n    Of that colossal Wreck, boundless and bare\n    The lone and level sands stretch far away\n    ",
    }
    # fmt: off
    _UpperCAmelCase			:int							=				[
        1864, 1536, 1213, 1870, 1357, 1536, 519, 880, 1323, 789, 1082, 534,
        1000, 1445, 1105, 1130, 967, 515, 1434, 1620, 534, 1495, 283, 1445,
        333, 1307, 539, 1631, 1528, 375, 1434, 673, 627, 710, 778, 1883,
        1405, 1276, 1455, 1228
    ]

    _UpperCAmelCase			:Optional[Any]							=				[
        1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653
    ]

    _UpperCAmelCase			:str							=				[
        1125, 1751, 697, 1776, 1141, 1476, 391, 697, 1125, 684, 867, 416,
        844, 1372, 1274, 717, 1274, 844, 1299, 1419, 697, 1370, 317, 1125,
        191, 1440, 1370, 1440, 1370, 282, 1621, 1370, 368, 349, 867, 1872,
        1262, 869, 1728, 747
    ]
    _UpperCAmelCase			:str							=				[
        416, 416, 1125, 1125, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416
    ]

    _UpperCAmelCase			:List[Any]							=				[
        1755, 842, 307, 1843, 1022, 1395, 234, 1554, 806, 739, 1022, 442,
        616, 556, 268, 1499, 933, 457, 1440, 1837, 755, 985, 308, 902,
        293, 1443, 1671, 1141, 1533, 555, 1562, 1061, 287, 417, 1022, 2008,
        1186, 1015, 1777, 268
    ]
    _UpperCAmelCase			:Tuple							=				[
        854, 842, 1353, 114, 1353, 842, 185, 842, 185, 114, 591, 842,
        185, 417, 185, 842, 307, 842, 591, 842, 185, 842, 307, 842,
        591, 842, 1353, 842, 185, 842, 591, 842, 591, 114, 591, 842,
        185, 842, 591, 89
    ]

    _UpperCAmelCase			:List[str]							=				[105_8304, 0, 78_6432, 7169, 507, 76, 27, 40, 30, 76]

    _UpperCAmelCase			:Union[str, Any]							=				[
        390, 1160, 1002, 1907, 1788, 1788, 1788, 1907, 1002, 1002, 1854, 1002,
        1002, 1002, 1002, 1002, 1002, 1160, 1160, 1606, 596, 596, 1160, 1002,
        1516, 596, 1002, 1002, 1002, 1907, 1788, 1788, 1788, 1854, 1788, 1907,
        1907, 1788, 596, 1626
    ]
    _UpperCAmelCase			:List[str]							=				[
        1236, 1668, 1484, 1920, 1848, 1409, 139, 864, 1828, 1272, 1599, 824,
        1672, 139, 555, 1484, 824, 1920, 555, 596, 1579, 1599, 1231, 1599,
        1637, 1407, 212, 824, 1599, 116, 1433, 824, 258, 1599, 1433, 1895,
        1063, 1433, 1433, 1599
    ]
    _UpperCAmelCase			:Any							=				[
        1684, 1873, 1119, 1189, 395, 611, 1901, 972, 890, 1337, 1392, 1927,
        96, 972, 672, 780, 1119, 890, 158, 771, 1073, 1927, 353, 1331,
        1269, 1459, 1333, 1645, 812, 1577, 1337, 606, 353, 981, 1466, 619,
        197, 391, 302, 1930
    ]
    _UpperCAmelCase			:List[str]							=				[
        390, 1160, 1002, 1907, 1788, 1788, 1788, 1907, 1002, 1002, 1854, 1002,
        1002, 1002, 1002, 1002, 1002, 1160, 1160, 1606, 596, 596, 1160, 1002,
        1516, 596, 1002, 1002, 1002, 1907, 1788, 1788, 1788, 1854, 1788, 1907,
        1907, 1788, 596, 1626
    ]
    _UpperCAmelCase			:List[Any]							=				[
        -0.0492, -0.0524, -0.0565, -0.0640, -0.0686, -0.0684, -0.0677, -0.0664,
        -0.0605, -0.0490, -0.0330, -0.0168, -0.0083, -0.0075, -0.0051, 0.0025,
        0.0136, 0.0261, 0.0386, 0.0497, 0.0580, 0.0599, 0.0583, 0.0614,
        0.0740, 0.0889, 0.1023, 0.1162, 0.1211, 0.1212, 0.1251, 0.1336,
        0.1502, 0.1686, 0.1883, 0.2148, 0.2363, 0.2458, 0.2507, 0.2531
    ]
    _UpperCAmelCase			:str							=				[
        0.0256, -0.0544, 0.1600, -0.0032, 0.1066, 0.0825, -0.0013, 0.3440,
        0.0210, 0.0412, -0.1777, -0.0892, -0.0164, 0.0285, -0.0613, -0.0617,
        -0.0137, -0.0201, -0.0175, 0.0215, -0.0627, 0.0520, -0.0730, 0.0970,
        -0.0100, 0.0442, -0.0586, 0.0207, -0.0015, -0.0082
    ]
    _UpperCAmelCase			:str							=				[
        0.0415, 0.0877, 0.0022, -0.0055, 0.0751, 0.0334, 0.0324, -0.0068,
        0.0011, 0.0017, -0.0676, 0.0655, -0.0143, 0.0399, 0.0303, 0.0743,
        -0.0168, -0.0394, -0.1113, 0.0124, 0.0442, 0.0267, -0.0003, -0.1536,
        -0.0116, -0.1837, -0.0180, -0.1026, -0.0777, -0.0456
    ]
    _UpperCAmelCase			:Optional[Any]							=				[
        76, 27, 40, 30, 76, 46, 44, 47, 40, 37, 38, 31, 45, 45, 76, 38, 31, 33,
        45, 76, 41, 32, 76, 45, 46, 41, 40, 31, 78, 76
    ]






    def _snake_case    (     self							):
       lowercase__:     Dict      =   JukeboxTokenizer.from_pretrained(self.model_id							)
       lowercase__:     str      =   tokenizer(**self.metas							)['''input_ids''']
       return tokens






    @slow
    def _snake_case    (     self							):
       lowercase__:     List[Any]      =   JukeboxModel.from_pretrained(self.model_id					,    min_duration=0							).eval()
       lowercase__:     Tuple      =   self.prepare_inputs()

       set_seed(0							)
       lowercase__:     Tuple      =   [torch.zeros(1					,    0					,    dtype=torch.long							).cpu() for _ in range(3							)]
       lowercase__:     Optional[Any]      =   model._sample(_UpperCAmelCase					,    _UpperCAmelCase					,    [0]					,    sample_length=40 * model.priors[0].raw_to_tokens					,    save_results=_UpperCAmelCase							)
       self.assertIn(zs[0][0].detach().cpu().tolist()					,    [self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]							)

       set_seed(0							)
       lowercase__:     str      =   model._sample(_UpperCAmelCase					,    _UpperCAmelCase					,    [1]					,    sample_length=40 * model.priors[1].raw_to_tokens					,    save_results=_UpperCAmelCase							)
       self.assertIn(zs[1][0].detach().cpu().tolist()					,    [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]							)

       set_seed(0							)
       lowercase__:     Any      =   model._sample(_UpperCAmelCase					,    _UpperCAmelCase					,    [2]					,    sample_length=40 * model.priors[2].raw_to_tokens					,    save_results=_UpperCAmelCase							)
       self.assertIn(zs[2][0].detach().cpu().tolist()					,    [self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]							)






    @slow
    def _snake_case    (     self							):
       lowercase__:     Dict      =   False
       lowercase__:     int      =   JukeboxModel.from_pretrained(self.model_id					,    min_duration=0							).eval()

       lowercase__:     List[str]      =   self.prepare_inputs()
       set_seed(0							)
       lowercase__:     Optional[Any]      =   [torch.zeros(1					,    0					,    dtype=torch.long							) for _ in range(3							)]

       lowercase__:     Union[str, Any]      =   model.priors[0]
       lowercase__:     Optional[int]      =   0
       lowercase__:     Any      =   top_prior.get_music_tokens_conds(_UpperCAmelCase					,    start=_UpperCAmelCase					,    end=start + top_prior.n_ctx							)
       lowercase__:     Optional[Any]      =   top_prior.get_metadata(labels[0].clone()					,    _UpperCAmelCase					,    1058304					,    0							)

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

       lowercase__,				lowercase__,				lowercase__:     Any      =   top_prior.get_cond(_UpperCAmelCase					,    _UpperCAmelCase							)
       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							)






    @slow
    def _snake_case    (     self							):
       lowercase__:     Any      =   False

       lowercase__:     int      =   JukeboxModel.from_pretrained(self.model_id					,    min_duration=0							).eval()
       set_seed(0							)
       lowercase__:     List[str]      =   torch.rand((1, 5120, 1)							)
       lowercase__:     Optional[Any]      =   list(self.prepare_inputs()							)

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

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

       lowercase__:     Optional[Any]      =   torch.cat((zs[1], torch.zeros(1					,    2048 - zs[1].shape[-1]							))					,    dim=-1							).long()
       lowercase__:     Union[str, Any]      =   [upper_a, upper_a, model.vqvae.encode(_UpperCAmelCase					,    start_level=0					,    bs_chunks=waveform.shape[0]							)[0]]
       lowercase__:     Optional[Any]      =   model._sample(
           _UpperCAmelCase					,    _UpperCAmelCase					,    sample_levels=[2]					,    save_results=_UpperCAmelCase					,    sample_length=40 * model.priors[2].raw_to_tokens							)
       torch.testing.assert_allclose(zs[2][0][:40].cpu()					,    torch.tensor(self.EXPECTED_PRIMED_2							)							)






    @slow
    def _snake_case    (     self							):
       lowercase__:     int      =   JukeboxModel.from_pretrained(self.model_id					,    min_duration=0							).eval()
       set_seed(0							)
       lowercase__:     Optional[int]      =   torch.rand((1, 5120, 1)							)
       with torch.no_grad():
          lowercase__:     Optional[Any]      =   model.vqvae.encode(_UpperCAmelCase					,    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():
          lowercase__:     Tuple      =   model.vqvae.decode(_UpperCAmelCase					,    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							)



@require_torch
class 					UpperCAmelCase		(unittest.TestCase						):



    """simple docstring"""

    _UpperCAmelCase			:List[Any]							=				(JukeboxModel,) if is_torch_available() else ()
    _UpperCAmelCase			:List[Any]							=				"openai/jukebox-5b-lyrics"
    _UpperCAmelCase			:Tuple							=				{
        "artist": "Zac Brown Band",
        "genres": "Country",
        "lyrics": "I met a traveller from an antique land,\n    Who said \"Two vast and trunkless legs of stone\n    Stand in the desert. . . . Near them, on the sand,\n    Half sunk a shattered visage lies, whose frown,\n    And wrinkled lip, and sneer of cold command,\n    Tell that its sculptor well those passions read\n    Which yet survive, stamped on these lifeless things,\n    The hand that mocked them, and the heart that fed;\n    And on the pedestal, these words appear:\n    My name is Ozymandias, King of Kings;\n    Look on my Works, ye Mighty, and despair!\n    Nothing beside remains. Round the decay\n    Of that colossal Wreck, boundless and bare\n    The lone and level sands stretch far away\n    ",
    }

    # fmt: off
    _UpperCAmelCase			:Dict							=				[
        1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        1489, 1489, 1489, 1489, 1150, 1853, 1509, 1150, 1357, 1509, 6, 1272
    ]
    _UpperCAmelCase			:Union[str, Any]							=				[
        1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
    ]

    _UpperCAmelCase			:Dict							=				[
        1125, 416, 1125, 1125, 1125, 1125, 1125, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416
    ]
    _UpperCAmelCase			:int							=				[
        416, 416, 1125, 1125, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416
    ]

    _UpperCAmelCase			:Any							=				[
        1755, 1061, 234, 1755, 1061, 1755, 185, 290, 307, 307, 616, 616,
        616, 616, 616, 616, 307, 290, 417, 1755, 234, 1755, 185, 290,
        290, 290, 307, 616, 616, 616, 616, 616, 290, 234, 234, 1755,
        234, 234, 1755, 234, 185, 185, 307, 616, 616, 616, 616, 290,
        1755, 1755, 1755, 234, 234, 1755, 1572, 290, 307, 616, 34, 616
    ]
    _UpperCAmelCase			:List[Any]							=				[
        854, 842, 1353, 114, 1353, 842, 185, 842, 185, 114, 591, 842, 185,
        417, 185, 842, 307, 842, 591, 842, 185, 842, 185, 842, 591, 842,
        1353, 842, 185, 842, 591, 842, 591, 114, 591, 842, 185, 842, 591,
        89, 591, 842, 591, 842, 591, 417, 1372, 842, 1372, 842, 34, 842,
        185, 89, 591, 842, 185, 842, 591, 632
    ]

    _UpperCAmelCase			:Tuple							=				[
        1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
    ]
    _UpperCAmelCase			:Tuple							=				[
        1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
        653, 653, 653, 653, 653, 653, 653, 1853, 1177, 1536, 1228,
        710, 475, 1489, 1229, 1224, 231, 1224, 252, 1434, 653, 475,
        1106, 1877, 1599, 1228, 1600, 1683, 1182, 1853, 475, 1864,
        252, 1229, 1434, 2001
    ]

    _UpperCAmelCase			:List[Any]							=				[
        1125, 1125, 416, 1125, 1125, 416, 1125, 1125, 416, 416, 1125, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
        416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416
    ]
    _UpperCAmelCase			:List[Any]							=				[
        491, 1755, 34, 1613, 1755, 417, 992, 1613, 222, 842, 1353, 1613,
        844, 632, 185, 1613, 844, 632, 185, 1613, 185, 842, 677, 1613,
        185, 114, 1353, 1613, 307, 89, 844, 1613, 307, 1332, 234, 1979,
        307, 89, 1353, 616, 34, 842, 185, 842, 34, 842, 185, 842,
        307, 114, 185, 89, 34, 1268, 185, 89, 34, 842, 185, 89
    ]






    def _snake_case    (     self					,    _UpperCAmelCase							):
       lowercase__:     Dict      =   JukeboxTokenizer.from_pretrained(_UpperCAmelCase							)
       lowercase__:     Optional[Any]      =   tokenizer(**self.metas							)['''input_ids''']
       return tokens






    @slow
    def _snake_case    (     self							):
       lowercase__:     Any      =   JukeboxModel.from_pretrained(self.model_id					,    min_duration=0							).eval()
       lowercase__:     int      =   self.prepare_inputs(self.model_id							)

       set_seed(0							)
       lowercase__:     Optional[int]      =   [torch.zeros(1					,    0					,    dtype=torch.long							).cpu() for _ in range(3							)]
       lowercase__:     Optional[int]      =   model._sample(_UpperCAmelCase					,    _UpperCAmelCase					,    [0]					,    sample_length=60 * model.priors[0].raw_to_tokens					,    save_results=_UpperCAmelCase							)
       self.assertIn(zs[0][0].detach().cpu().tolist()					,    [self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]							)

       set_seed(0							)
       lowercase__:     Optional[Any]      =   model._sample(_UpperCAmelCase					,    _UpperCAmelCase					,    [1]					,    sample_length=60 * model.priors[1].raw_to_tokens					,    save_results=_UpperCAmelCase							)
       self.assertIn(zs[1][0].detach().cpu().tolist()					,    [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]							)

       set_seed(0							)
       lowercase__:     Optional[Any]      =   model._sample(_UpperCAmelCase					,    _UpperCAmelCase					,    [2]					,    sample_length=60 * model.priors[2].raw_to_tokens					,    save_results=_UpperCAmelCase							)
       self.assertIn(zs[2][0].detach().cpu().tolist()					,    [self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]							)






    @slow
    @skip('''Not enough GPU memory on CI runners'''							)
    def _snake_case    (     self							):
       lowercase__:     List[str]      =   JukeboxModel.from_pretrained(self.model_id					,    min_duration=0							).eval()
       lowercase__:     Optional[Any]      =   [i.cuda() for i in self.prepare_inputs(self.model_id							)]

       set_seed(0							)
       model.priors[0].cuda()
       lowercase__:     int      =   [torch.zeros(1					,    0					,    dtype=torch.long							).cuda() for _ in range(3							)]
       lowercase__:     Dict      =   model._sample(_UpperCAmelCase					,    _UpperCAmelCase					,    [0]					,    sample_length=60 * model.priors[0].raw_to_tokens					,    save_results=_UpperCAmelCase							)
       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()
       lowercase__:     Optional[Any]      =   model._sample(_UpperCAmelCase					,    _UpperCAmelCase					,    [1]					,    sample_length=60 * model.priors[1].raw_to_tokens					,    save_results=_UpperCAmelCase							)
       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()
       lowercase__:     List[str]      =   model._sample(_UpperCAmelCase					,    _UpperCAmelCase					,    [2]					,    sample_length=60 * model.priors[2].raw_to_tokens					,    save_results=_UpperCAmelCase							)
       torch.testing.assert_allclose(zs[2][0].cpu()					,    torch.tensor(self.EXPECTED_GPU_OUTPUTS_0							)							)






    @slow
    def _snake_case    (     self							):
       lowercase__:     Tuple      =   '''ArthurZ/jukebox_prior_0'''
       lowercase__:     Any      =   JukeboxPrior.from_pretrained(_UpperCAmelCase					,    min_duration=0							).eval().half().to('''cuda'''							)

       lowercase__:     Tuple      =   self.prepare_inputs(_UpperCAmelCase							)[0].cuda()
       lowercase__:     Tuple      =   model.get_metadata(_UpperCAmelCase					,    0					,    7680					,    0							)
       set_seed(0							)
       lowercase__:     List[Any]      =   model.sample(1					,    metadata=_UpperCAmelCase					,    sample_tokens=60							)
       self.assertIn(outputs[0].cpu().tolist()					,    [self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2]							)

"""simple docstring"""






def SCREAMING_SNAKE_CASE__    (     __UpperCAmelCase = 1_0_0_0				)						->						int:
   lowercase__:     Union[str, Any]      =   0
   lowercase__:     Any      =   0
   while 1:
      num += 3
      if num >= n:
         break
      total += num
      num += 2
      if num >= n:
         break
      total += num
      num += 1
      if num >= n:
         break
      total += num
      num += 3
      if num >= n:
         break
      total += num
      num += 1
      if num >= n:
         break
      total += num
      num += 2
      if num >= n:
         break
      total += num
      num += 3
      if num >= n:
         break
      total += num
   return total


if __name__ == "__main__":
       print(f'''{solution() = }''')
