"""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 			_lowercase		(   unittest.TestCase  ):
				lowerCamelCase__						=       (JukeboxModel,) if is_torch_available() else ()
				lowerCamelCase__						=       'openai/jukebox-1b-lyrics'
				lowerCamelCase__						=       {
				    '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
				lowerCamelCase__						=       [
				    1_8_6_4, 1_5_3_6, 1_2_1_3, 1_8_7_0, 1_3_5_7, 1_5_3_6, 5_1_9, 8_8_0, 1_3_2_3, 7_8_9, 1_0_8_2, 5_3_4,
				    1_0_0_0, 1_4_4_5, 1_1_0_5, 1_1_3_0, 9_6_7, 5_1_5, 1_4_3_4, 1_6_2_0, 5_3_4, 1_4_9_5, 2_8_3, 1_4_4_5,
				    3_3_3, 1_3_0_7, 5_3_9, 1_6_3_1, 1_5_2_8, 3_7_5, 1_4_3_4, 6_7_3, 6_2_7, 7_1_0, 7_7_8, 1_8_8_3,
				    1_4_0_5, 1_2_7_6, 1_4_5_5, 1_2_2_8
				]

				lowerCamelCase__						=       [
				    1_4_8_9, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3
				]

				lowerCamelCase__						=       [
				    1_1_2_5, 1_7_5_1, 6_9_7, 1_7_7_6, 1_1_4_1, 1_4_7_6, 3_9_1, 6_9_7, 1_1_2_5, 6_8_4, 8_6_7, 4_1_6,
				    8_4_4, 1_3_7_2, 1_2_7_4, 7_1_7, 1_2_7_4, 8_4_4, 1_2_9_9, 1_4_1_9, 6_9_7, 1_3_7_0, 3_1_7, 1_1_2_5,
				    1_9_1, 1_4_4_0, 1_3_7_0, 1_4_4_0, 1_3_7_0, 2_8_2, 1_6_2_1, 1_3_7_0, 3_6_8, 3_4_9, 8_6_7, 1_8_7_2,
				    1_2_6_2, 8_6_9, 1_7_2_8, 7_4_7
				]
				lowerCamelCase__						=       [
				    4_1_6, 4_1_6, 1_1_2_5, 1_1_2_5, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6
				]

				lowerCamelCase__						=       [
				    1_7_5_5, 8_4_2, 3_0_7, 1_8_4_3, 1_0_2_2, 1_3_9_5, 2_3_4, 1_5_5_4, 8_0_6, 7_3_9, 1_0_2_2, 4_4_2,
				    6_1_6, 5_5_6, 2_6_8, 1_4_9_9, 9_3_3, 4_5_7, 1_4_4_0, 1_8_3_7, 7_5_5, 9_8_5, 3_0_8, 9_0_2,
				    2_9_3, 1_4_4_3, 1_6_7_1, 1_1_4_1, 1_5_3_3, 5_5_5, 1_5_6_2, 1_0_6_1, 2_8_7, 4_1_7, 1_0_2_2, 2_0_0_8,
				    1_1_8_6, 1_0_1_5, 1_7_7_7, 2_6_8
				]
				lowerCamelCase__						=       [
				    8_5_4, 8_4_2, 1_3_5_3, 1_1_4, 1_3_5_3, 8_4_2, 1_8_5, 8_4_2, 1_8_5, 1_1_4, 5_9_1, 8_4_2,
				    1_8_5, 4_1_7, 1_8_5, 8_4_2, 3_0_7, 8_4_2, 5_9_1, 8_4_2, 1_8_5, 8_4_2, 3_0_7, 8_4_2,
				    5_9_1, 8_4_2, 1_3_5_3, 8_4_2, 1_8_5, 8_4_2, 5_9_1, 8_4_2, 5_9_1, 1_1_4, 5_9_1, 8_4_2,
				    1_8_5, 8_4_2, 5_9_1, 8_9
				]

				lowerCamelCase__						=       [1_0_5_8_3_0_4, 0, 7_8_6_4_3_2, 7_1_6_9, 5_0_7, 7_6, 2_7, 4_0, 3_0, 7_6]

				lowerCamelCase__						=       [
				    3_9_0, 1_1_6_0, 1_0_0_2, 1_9_0_7, 1_7_8_8, 1_7_8_8, 1_7_8_8, 1_9_0_7, 1_0_0_2, 1_0_0_2, 1_8_5_4, 1_0_0_2,
				    1_0_0_2, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_1_6_0, 1_1_6_0, 1_6_0_6, 5_9_6, 5_9_6, 1_1_6_0, 1_0_0_2,
				    1_5_1_6, 5_9_6, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_9_0_7, 1_7_8_8, 1_7_8_8, 1_7_8_8, 1_8_5_4, 1_7_8_8, 1_9_0_7,
				    1_9_0_7, 1_7_8_8, 5_9_6, 1_6_2_6
				]
				lowerCamelCase__						=       [
				    1_2_3_6, 1_6_6_8, 1_4_8_4, 1_9_2_0, 1_8_4_8, 1_4_0_9, 1_3_9, 8_6_4, 1_8_2_8, 1_2_7_2, 1_5_9_9, 8_2_4,
				    1_6_7_2, 1_3_9, 5_5_5, 1_4_8_4, 8_2_4, 1_9_2_0, 5_5_5, 5_9_6, 1_5_7_9, 1_5_9_9, 1_2_3_1, 1_5_9_9,
				    1_6_3_7, 1_4_0_7, 2_1_2, 8_2_4, 1_5_9_9, 1_1_6, 1_4_3_3, 8_2_4, 2_5_8, 1_5_9_9, 1_4_3_3, 1_8_9_5,
				    1_0_6_3, 1_4_3_3, 1_4_3_3, 1_5_9_9
				]
				lowerCamelCase__						=       [
				    1_6_8_4, 1_8_7_3, 1_1_1_9, 1_1_8_9, 3_9_5, 6_1_1, 1_9_0_1, 9_7_2, 8_9_0, 1_3_3_7, 1_3_9_2, 1_9_2_7,
				    9_6, 9_7_2, 6_7_2, 7_8_0, 1_1_1_9, 8_9_0, 1_5_8, 7_7_1, 1_0_7_3, 1_9_2_7, 3_5_3, 1_3_3_1,
				    1_2_6_9, 1_4_5_9, 1_3_3_3, 1_6_4_5, 8_1_2, 1_5_7_7, 1_3_3_7, 6_0_6, 3_5_3, 9_8_1, 1_4_6_6, 6_1_9,
				    1_9_7, 3_9_1, 3_0_2, 1_9_3_0
				]
				lowerCamelCase__						=       [
				    3_9_0, 1_1_6_0, 1_0_0_2, 1_9_0_7, 1_7_8_8, 1_7_8_8, 1_7_8_8, 1_9_0_7, 1_0_0_2, 1_0_0_2, 1_8_5_4, 1_0_0_2,
				    1_0_0_2, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_1_6_0, 1_1_6_0, 1_6_0_6, 5_9_6, 5_9_6, 1_1_6_0, 1_0_0_2,
				    1_5_1_6, 5_9_6, 1_0_0_2, 1_0_0_2, 1_0_0_2, 1_9_0_7, 1_7_8_8, 1_7_8_8, 1_7_8_8, 1_8_5_4, 1_7_8_8, 1_9_0_7,
				    1_9_0_7, 1_7_8_8, 5_9_6, 1_6_2_6
				]
				lowerCamelCase__						=       [
				    -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
				]
				lowerCamelCase__						=       [
				    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
				]
				lowerCamelCase__						=       [
				    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
				]
				lowerCamelCase__						=       [
				    7_6, 2_7, 4_0, 3_0, 7_6, 4_6, 4_4, 4_7, 4_0, 3_7, 3_8, 3_1, 4_5, 4_5, 7_6, 3_8, 3_1, 3_3,
				    4_5, 7_6, 4_1, 3_2, 7_6, 4_5, 4_6, 4_1, 4_0, 3_1, 7_8, 7_6
				]

				def        A (     self       :    int     ):


										'''simple docstring'''
										UpperCAmelCase__						=		JukeboxTokenizer.from_pretrained(self.model_id     )
										UpperCAmelCase__						=		tokenizer(**self.metas     )["""input_ids"""]
										return tokens

				@slow
				def        A (     self       :    Dict     ):


										'''simple docstring'''
										UpperCAmelCase__						=		JukeboxModel.from_pretrained(self.model_id       ,    min_duration=0     ).eval()
										UpperCAmelCase__						=		self.prepare_inputs()

										set_seed(0     )
										UpperCAmelCase__						=		[torch.zeros(1       ,    0       ,    dtype=torch.long     ).cpu() for _ in range(3     )]
										UpperCAmelCase__						=		model._sample(_lowercase       ,    _lowercase       ,    [0]       ,    sample_length=4_0 * model.priors[0].raw_to_tokens       ,    save_results=_lowercase     )
										self.assertIn(zs[0][0].detach().cpu().tolist()       ,    [self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]     )

										set_seed(0     )
										UpperCAmelCase__						=		model._sample(_lowercase       ,    _lowercase       ,    [1]       ,    sample_length=4_0 * model.priors[1].raw_to_tokens       ,    save_results=_lowercase     )
										self.assertIn(zs[1][0].detach().cpu().tolist()       ,    [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]     )

										set_seed(0     )
										UpperCAmelCase__						=		model._sample(_lowercase       ,    _lowercase       ,    [2]       ,    sample_length=4_0 * model.priors[2].raw_to_tokens       ,    save_results=_lowercase     )
										self.assertIn(zs[2][0].detach().cpu().tolist()       ,    [self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]     )

				@slow
				def        A (     self       :    List[str]     ):


										'''simple docstring'''
										UpperCAmelCase__						=		False
										UpperCAmelCase__						=		JukeboxModel.from_pretrained(self.model_id       ,    min_duration=0     ).eval()

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

										UpperCAmelCase__						=		model.priors[0]
										UpperCAmelCase__						=		0
										UpperCAmelCase__						=		top_prior.get_music_tokens_conds(_lowercase       ,    start=_lowercase       ,    end=start + top_prior.n_ctx     )
										UpperCAmelCase__						=		top_prior.get_metadata(labels[0].clone()       ,    _lowercase       ,    1_0_5_8_3_0_4       ,    0     )

										self.assertIsNone(_lowercase     )
										self.assertListEqual(metadata.numpy()[0][:1_0].tolist()       ,    self.EXPECTED_Y_COND     )

										UpperCAmelCase__	,      UpperCAmelCase__	,      UpperCAmelCase__						=		top_prior.get_cond(_lowercase       ,    _lowercase     )
										torch.testing.assert_allclose(
										    audio_conditioning[0][0][:3_0].detach()       ,    torch.tensor(self.EXPECTED_AUDIO_COND     )       ,    atol=1e-4       ,    rtol=1e-4     )
										torch.testing.assert_allclose(
										    metadata_conditioning[0][0][:3_0].detach()       ,    torch.tensor(self.EXPECTED_META_COND     )       ,    atol=1e-4       ,    rtol=1e-4     )
										torch.testing.assert_allclose(
										    lyric_tokens[0, :3_0].detach()       ,    torch.tensor(self.EXPECTED_LYRIC_COND     )       ,    atol=1e-4       ,    rtol=1e-4     )

				@slow
				def        A (     self       :    Union[str, Any]     ):


										'''simple docstring'''
										UpperCAmelCase__						=		False

										UpperCAmelCase__						=		JukeboxModel.from_pretrained(self.model_id       ,    min_duration=0     ).eval()
										set_seed(0     )
										UpperCAmelCase__						=		torch.rand((1, 5_1_2_0, 1)     )
										UpperCAmelCase__						=		list(self.prepare_inputs()     )

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

										UpperCAmelCase__						=		torch.cat((zs[0], torch.zeros(1       ,    2_0_4_8 - zs[0].shape[-1]     ))       ,    dim=-1     ).long()
										UpperCAmelCase__						=		[upper_a, model.vqvae.encode(_lowercase       ,    start_level=1       ,    bs_chunks=waveform.shape[0]     )[0], None]
										UpperCAmelCase__						=		model._sample(
										    _lowercase       ,    _lowercase       ,    sample_levels=[1]       ,    save_results=_lowercase       ,    sample_length=4_0 * model.priors[1].raw_to_tokens     )
										torch.testing.assert_allclose(zs[1][0][:4_0]       ,    torch.tensor(self.EXPECTED_PRIMED_1     )     )

										UpperCAmelCase__						=		torch.cat((zs[1], torch.zeros(1       ,    2_0_4_8 - zs[1].shape[-1]     ))       ,    dim=-1     ).long()
										UpperCAmelCase__						=		[upper_a, upper_a, model.vqvae.encode(_lowercase       ,    start_level=0       ,    bs_chunks=waveform.shape[0]     )[0]]
										UpperCAmelCase__						=		model._sample(
										    _lowercase       ,    _lowercase       ,    sample_levels=[2]       ,    save_results=_lowercase       ,    sample_length=4_0 * model.priors[2].raw_to_tokens     )
										torch.testing.assert_allclose(zs[2][0][:4_0].cpu()       ,    torch.tensor(self.EXPECTED_PRIMED_2     )     )

				@slow
				def        A (     self       :    Optional[int]     ):


										'''simple docstring'''
										UpperCAmelCase__						=		JukeboxModel.from_pretrained(self.model_id       ,    min_duration=0     ).eval()
										set_seed(0     )
										UpperCAmelCase__						=		torch.rand((1, 5_1_2_0, 1)     )
										with torch.no_grad():
																UpperCAmelCase__						=		model.vqvae.encode(_lowercase       ,    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():
																UpperCAmelCase__						=		model.vqvae.decode(_lowercase       ,    start_level=2       ,    bs_chunks=x.shape[0]     )
										torch.testing.assert_allclose(x[0, :4_0, 0]       ,    torch.tensor(self.EXPECTED_VQVAE_DECODE     )       ,    atol=1e-4       ,    rtol=1e-4     )






@require_torch
class 			_lowercase		(   unittest.TestCase  ):
				lowerCamelCase__						=       (JukeboxModel,) if is_torch_available() else ()
				lowerCamelCase__						=       'openai/jukebox-5b-lyrics'
				lowerCamelCase__						=       {
				    '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
				lowerCamelCase__						=       [
				    1_4_8_9, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    1_4_8_9, 1_4_8_9, 1_4_8_9, 1_4_8_9, 1_1_5_0, 1_8_5_3, 1_5_0_9, 1_1_5_0, 1_3_5_7, 1_5_0_9, 6, 1_2_7_2
				]
				lowerCamelCase__						=       [
				    1_4_8_9, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3
				]

				lowerCamelCase__						=       [
				    1_1_2_5, 4_1_6, 1_1_2_5, 1_1_2_5, 1_1_2_5, 1_1_2_5, 1_1_2_5, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6
				]
				lowerCamelCase__						=       [
				    4_1_6, 4_1_6, 1_1_2_5, 1_1_2_5, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6
				]

				lowerCamelCase__						=       [
				    1_7_5_5, 1_0_6_1, 2_3_4, 1_7_5_5, 1_0_6_1, 1_7_5_5, 1_8_5, 2_9_0, 3_0_7, 3_0_7, 6_1_6, 6_1_6,
				    6_1_6, 6_1_6, 6_1_6, 6_1_6, 3_0_7, 2_9_0, 4_1_7, 1_7_5_5, 2_3_4, 1_7_5_5, 1_8_5, 2_9_0,
				    2_9_0, 2_9_0, 3_0_7, 6_1_6, 6_1_6, 6_1_6, 6_1_6, 6_1_6, 2_9_0, 2_3_4, 2_3_4, 1_7_5_5,
				    2_3_4, 2_3_4, 1_7_5_5, 2_3_4, 1_8_5, 1_8_5, 3_0_7, 6_1_6, 6_1_6, 6_1_6, 6_1_6, 2_9_0,
				    1_7_5_5, 1_7_5_5, 1_7_5_5, 2_3_4, 2_3_4, 1_7_5_5, 1_5_7_2, 2_9_0, 3_0_7, 6_1_6, 3_4, 6_1_6
				]
				lowerCamelCase__						=       [
				    8_5_4, 8_4_2, 1_3_5_3, 1_1_4, 1_3_5_3, 8_4_2, 1_8_5, 8_4_2, 1_8_5, 1_1_4, 5_9_1, 8_4_2, 1_8_5,
				    4_1_7, 1_8_5, 8_4_2, 3_0_7, 8_4_2, 5_9_1, 8_4_2, 1_8_5, 8_4_2, 1_8_5, 8_4_2, 5_9_1, 8_4_2,
				    1_3_5_3, 8_4_2, 1_8_5, 8_4_2, 5_9_1, 8_4_2, 5_9_1, 1_1_4, 5_9_1, 8_4_2, 1_8_5, 8_4_2, 5_9_1,
				    8_9, 5_9_1, 8_4_2, 5_9_1, 8_4_2, 5_9_1, 4_1_7, 1_3_7_2, 8_4_2, 1_3_7_2, 8_4_2, 3_4, 8_4_2,
				    1_8_5, 8_9, 5_9_1, 8_4_2, 1_8_5, 8_4_2, 5_9_1, 6_3_2
				]

				lowerCamelCase__						=       [
				    1_4_8_9, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3
				]
				lowerCamelCase__						=       [
				    1_4_8_9, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3,
				    6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 6_5_3, 1_8_5_3, 1_1_7_7, 1_5_3_6, 1_2_2_8,
				    7_1_0, 4_7_5, 1_4_8_9, 1_2_2_9, 1_2_2_4, 2_3_1, 1_2_2_4, 2_5_2, 1_4_3_4, 6_5_3, 4_7_5,
				    1_1_0_6, 1_8_7_7, 1_5_9_9, 1_2_2_8, 1_6_0_0, 1_6_8_3, 1_1_8_2, 1_8_5_3, 4_7_5, 1_8_6_4,
				    2_5_2, 1_2_2_9, 1_4_3_4, 2_0_0_1
				]

				lowerCamelCase__						=       [
				    1_1_2_5, 1_1_2_5, 4_1_6, 1_1_2_5, 1_1_2_5, 4_1_6, 1_1_2_5, 1_1_2_5, 4_1_6, 4_1_6, 1_1_2_5, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6,
				    4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6, 4_1_6
				]
				lowerCamelCase__						=       [
				    4_9_1, 1_7_5_5, 3_4, 1_6_1_3, 1_7_5_5, 4_1_7, 9_9_2, 1_6_1_3, 2_2_2, 8_4_2, 1_3_5_3, 1_6_1_3,
				    8_4_4, 6_3_2, 1_8_5, 1_6_1_3, 8_4_4, 6_3_2, 1_8_5, 1_6_1_3, 1_8_5, 8_4_2, 6_7_7, 1_6_1_3,
				    1_8_5, 1_1_4, 1_3_5_3, 1_6_1_3, 3_0_7, 8_9, 8_4_4, 1_6_1_3, 3_0_7, 1_3_3_2, 2_3_4, 1_9_7_9,
				    3_0_7, 8_9, 1_3_5_3, 6_1_6, 3_4, 8_4_2, 1_8_5, 8_4_2, 3_4, 8_4_2, 1_8_5, 8_4_2,
				    3_0_7, 1_1_4, 1_8_5, 8_9, 3_4, 1_2_6_8, 1_8_5, 8_9, 3_4, 8_4_2, 1_8_5, 8_9
				]

				def        A (     self       :    Dict       ,    _lowercase       :    Optional[int]     ):


										'''simple docstring'''
										UpperCAmelCase__						=		JukeboxTokenizer.from_pretrained(_lowercase     )
										UpperCAmelCase__						=		tokenizer(**self.metas     )["""input_ids"""]
										return tokens

				@slow
				def        A (     self       :    int     ):


										'''simple docstring'''
										UpperCAmelCase__						=		JukeboxModel.from_pretrained(self.model_id       ,    min_duration=0     ).eval()
										UpperCAmelCase__						=		self.prepare_inputs(self.model_id     )

										set_seed(0     )
										UpperCAmelCase__						=		[torch.zeros(1       ,    0       ,    dtype=torch.long     ).cpu() for _ in range(3     )]
										UpperCAmelCase__						=		model._sample(_lowercase       ,    _lowercase       ,    [0]       ,    sample_length=6_0 * model.priors[0].raw_to_tokens       ,    save_results=_lowercase     )
										self.assertIn(zs[0][0].detach().cpu().tolist()       ,    [self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]     )

										set_seed(0     )
										UpperCAmelCase__						=		model._sample(_lowercase       ,    _lowercase       ,    [1]       ,    sample_length=6_0 * model.priors[1].raw_to_tokens       ,    save_results=_lowercase     )
										self.assertIn(zs[1][0].detach().cpu().tolist()       ,    [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]     )

										set_seed(0     )
										UpperCAmelCase__						=		model._sample(_lowercase       ,    _lowercase       ,    [2]       ,    sample_length=6_0 * model.priors[2].raw_to_tokens       ,    save_results=_lowercase     )
										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        A (     self       :    List[str]     ):


										'''simple docstring'''
										UpperCAmelCase__						=		JukeboxModel.from_pretrained(self.model_id       ,    min_duration=0     ).eval()
										UpperCAmelCase__						=		[i.cuda() for i in self.prepare_inputs(self.model_id     )]

										set_seed(0     )
										model.priors[0].cuda()
										UpperCAmelCase__						=		[torch.zeros(1       ,    0       ,    dtype=torch.long     ).cuda() for _ in range(3     )]
										UpperCAmelCase__						=		model._sample(_lowercase       ,    _lowercase       ,    [0]       ,    sample_length=6_0 * model.priors[0].raw_to_tokens       ,    save_results=_lowercase     )
										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()
										UpperCAmelCase__						=		model._sample(_lowercase       ,    _lowercase       ,    [1]       ,    sample_length=6_0 * model.priors[1].raw_to_tokens       ,    save_results=_lowercase     )
										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()
										UpperCAmelCase__						=		model._sample(_lowercase       ,    _lowercase       ,    [2]       ,    sample_length=6_0 * model.priors[2].raw_to_tokens       ,    save_results=_lowercase     )
										torch.testing.assert_allclose(zs[2][0].cpu()       ,    torch.tensor(self.EXPECTED_GPU_OUTPUTS_0     )     )

				@slow
				def        A (     self       :    Optional[Any]     ):


										'''simple docstring'''
										UpperCAmelCase__						=		"""ArthurZ/jukebox_prior_0"""
										UpperCAmelCase__						=		JukeboxPrior.from_pretrained(_lowercase       ,    min_duration=0     ).eval().half().to("""cuda"""     )

										UpperCAmelCase__						=		self.prepare_inputs(_lowercase     )[0].cuda()
										UpperCAmelCase__						=		model.get_metadata(_lowercase       ,    0       ,    7_6_8_0       ,    0     )
										set_seed(0     )
										UpperCAmelCase__						=		model.sample(1       ,    metadata=_lowercase       ,    sample_tokens=6_0     )
										self.assertIn(outputs[0].cpu().tolist()       ,    [self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2]     )


"""simple docstring"""







import math
from timeit import timeit


def   a__      (		UpperCAmelCase__       :			int			)->	int:
						UpperCAmelCase__						=		0
						UpperCAmelCase__						=		abs(UpperCAmelCase__			)
						while True:
												UpperCAmelCase__						=		n // 1_0
												digits += 1
												if n == 0:
																		break
						return digits


def   a__      (		UpperCAmelCase__       :			int			)->	int:
						return 1 if n == 0 else math.floor(math.log(abs(UpperCAmelCase__			)		,			1_0			) + 1			)


def   a__      (		UpperCAmelCase__       :			int			)->	int:
						return len(str(abs(UpperCAmelCase__			)			)			)


def   a__      (		)->	None:
						from collections.abc import Callable

						def benchmark_a_function(UpperCAmelCase__       :			Callable		,			UpperCAmelCase__       :			int			) -> None:
												UpperCAmelCase__						=		f'{func.__name__}({value})'
												UpperCAmelCase__						=		timeit(f'__main__.{call}'		,			setup="""import __main__"""			)
												print(f'{call}: {func(UpperCAmelCase__			)} -- {timing} seconds'			)

						for value in (2_6_2_1_4_4, 1_1_2_5_8_9_9_9_0_6_8_4_2_6_2_4, 1_2_6_7_6_5_0_6_0_0_2_2_8_2_2_9_4_0_1_4_9_6_7_0_3_2_0_5_3_7_6):
												for func in (num_digits, num_digits_fast, num_digits_faster):
																		benchmark_a_function(UpperCAmelCase__		,			UpperCAmelCase__			)
												print()


if __name__ == "__main__":
						import doctest

						doctest.testmod()
						benchmark()

