




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"""






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

	a     : Any	  =					[
	    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
	]

	a     : Union[str, Any]	  =					[
	    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
	]
	a     : List[Any]	  =					[
	    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
	]

	a     : int	  =					[
	    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
	]
	a     : Tuple	  =					[
	    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
	]

	a     : str	  =					[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]

	a     : List[Any]	  =					[
	    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
	]
	a     : Dict	  =					[
	    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
	]
	a     : List[str]	  =					[
	    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
	]
	a     : Optional[Any]	  =					[
	    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
	]
	a     : Dict	  =					[
	    -0.0_492, -0.0_524, -0.0_565, -0.0_640, -0.0_686, -0.0_684, -0.0_677, -0.0_664,
	    -0.0_605, -0.0_490, -0.0_330, -0.0_168, -0.0_083, -0.0_075, -0.0_051, 0.0_025,
	    0.0_136, 0.0_261, 0.0_386, 0.0_497, 0.0_580, 0.0_599, 0.0_583, 0.0_614,
	    0.0_740, 0.0_889, 0.1_023, 0.1_162, 0.1_211, 0.1_212, 0.1_251, 0.1_336,
	    0.1_502, 0.1_686, 0.1_883, 0.2_148, 0.2_363, 0.2_458, 0.2_507, 0.2_531
	]
	a     : Union[str, Any]	  =					[
	    0.0_256, -0.0_544, 0.1_600, -0.0_032, 0.1_066, 0.0_825, -0.0_013, 0.3_440,
	    0.0_210, 0.0_412, -0.1_777, -0.0_892, -0.0_164, 0.0_285, -0.0_613, -0.0_617,
	    -0.0_137, -0.0_201, -0.0_175, 0.0_215, -0.0_627, 0.0_520, -0.0_730, 0.0_970,
	    -0.0_100, 0.0_442, -0.0_586, 0.0_207, -0.0_015, -0.0_082
	]
	a     : Optional[int]	  =					[
	    0.0_415, 0.0_877, 0.0_022, -0.0_055, 0.0_751, 0.0_334, 0.0_324, -0.0_068,
	    0.0_011, 0.0_017, -0.0_676, 0.0_655, -0.0_143, 0.0_399, 0.0_303, 0.0_743,
	    -0.0_168, -0.0_394, -0.1_113, 0.0_124, 0.0_442, 0.0_267, -0.0_003, -0.1_536,
	    -0.0_116, -0.1_837, -0.0_180, -0.1_026, -0.0_777, -0.0_456
	]
	a     : List[Any]	  =					[
	    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    ) ->      Dict:





								"""simple docstring"""


								lowercase__        =							JukeboxTokenizer.from_pretrained(self.model_id    )
								lowercase__        =							tokenizer(**self.metas    )["input_ids"]
								return tokens

	@slow
	def    __a     (					self    ) ->      List[str]:





								"""simple docstring"""


								lowercase__        =							JukeboxModel.from_pretrained(self.model_id	,				min_duration=0    ).eval()
								lowercase__        =							self.prepare_inputs()

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

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

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

	@slow
	def    __a     (					self    ) ->      Optional[Any]:





								"""simple docstring"""


								lowercase__        =							False
								lowercase__        =							JukeboxModel.from_pretrained(self.model_id	,				min_duration=0    ).eval()

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

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

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

								lowercase__        =							top_prior.get_cond(lowerCamelCase	,				lowerCamelCase    )
								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    __a     (					self    ) ->      int:





								"""simple docstring"""


								lowercase__        =							False

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

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

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

								lowercase__        =							torch.cat((zs[1], torch.zeros(1	,				2048 - zs[1].shape[-1]    ))	,				dim=-1    ).long()
								lowercase__        =							[upper_a, upper_a, model.vqvae.encode(lowerCamelCase	,				start_level=0	,				bs_chunks=waveform.shape[0]    )[0]]
								lowercase__        =							model._sample(
								    lowerCamelCase	,				lowerCamelCase	,				sample_levels=[2]	,				save_results=lowerCamelCase	,				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    __a     (					self    ) ->      str:





								"""simple docstring"""


								lowercase__        =							JukeboxModel.from_pretrained(self.model_id	,				min_duration=0    ).eval()
								set_seed(0    )
								lowercase__        =							torch.rand((1, 5120, 1)    )
								with torch.no_grad():
															lowercase__        =							model.vqvae.encode(lowerCamelCase	,				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__        =							model.vqvae.decode(lowerCamelCase	,				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"""






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

	# fmt: off
	a     : int	  =					[
	    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
	]
	a     : Optional[int]	  =					[
	    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
	]

	a     : Union[str, Any]	  =					[
	    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
	]
	a     : Optional[int]	  =					[
	    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
	]

	a     : int	  =					[
	    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
	]
	a     : List[str]	  =					[
	    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
	]

	a     : Optional[Any]	  =					[
	    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
	]
	a     : List[str]	  =					[
	    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
	]

	a     : int	  =					[
	    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
	]
	a     : int	  =					[
	    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	,				lowerCamelCase    ) ->      str:





								"""simple docstring"""


								lowercase__        =							JukeboxTokenizer.from_pretrained(lowerCamelCase    )
								lowercase__        =							tokenizer(**self.metas    )["input_ids"]
								return tokens

	@slow
	def    __a     (					self    ) ->      Optional[int]:





								"""simple docstring"""


								lowercase__        =							JukeboxModel.from_pretrained(self.model_id	,				min_duration=0    ).eval()
								lowercase__        =							self.prepare_inputs(self.model_id    )

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

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

								set_seed(0    )
								lowercase__        =							model._sample(lowerCamelCase	,				lowerCamelCase	,				[2]	,				sample_length=60 * model.priors[2].raw_to_tokens	,				save_results=lowerCamelCase    )
								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    ) ->      Optional[Any]:





								"""simple docstring"""


								lowercase__        =							JukeboxModel.from_pretrained(self.model_id	,				min_duration=0    ).eval()
								lowercase__        =							[i.cuda() for i in self.prepare_inputs(self.model_id    )]

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







	@slow
	def    __a     (					self    ) ->      str:





								"""simple docstring"""


								lowercase__        =							"ArthurZ/jukebox_prior_0"
								lowercase__        =							JukeboxPrior.from_pretrained(lowerCamelCase	,				min_duration=0    ).eval().half().to("cuda"    )

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





def  snake_case_					(		)      ->			int:

							for a in range(3_00    ):
														for b in range(a + 1    ,4_00    ):
																					for c in range(b + 1    ,5_00    ):
																												if (a + b + c) == 10_00 and (a**2) + (b**2) == (c**2):
																																			return a * b * c

							return -1
def  snake_case_					(		)      ->			int:

							for a in range(3_00    ):
														for b in range(4_00    ):
																					lowercase__  :	Dict        =							10_00 - a - b
																					if a < b < c and (a**2) + (b**2) == (c**2):
																												return a * b * c

							return -1
def  snake_case_					(		)      ->			None:
							import timeit

							print(
							    timeit.timeit("solution()"    ,setup="from __main__ import solution"    ,number=10_00    )    )
							print(
							    timeit.timeit(
							        "solution_fast()"    ,setup="from __main__ import solution_fast"    ,number=10_00    )    )


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