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



       __lowercase   :   List[Any]  					=							(JukeboxModel,) if is_torch_available() else ()
       __lowercase   :   Optional[Any]  					=							'openai/jukebox-1b-lyrics'
       __lowercase   :   Any  					=							{
           '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
       __lowercase   :   Dict  					=							[
           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
       ]

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

       __lowercase   :   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
       ]
       __lowercase   :   List[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
       ]

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

       __lowercase   :   Union[str, Any]  					=							[1058304, 0, 786432, 7169, 507, 76, 27, 40, 30, 76]

       __lowercase   :   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
       ]
       __lowercase   :   int  					=							[
           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
       ]
       __lowercase   :   Tuple  					=							[
           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
       ]
       __lowercase   :   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
       ]
       __lowercase   :   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
       ]
       __lowercase   :   Union[str, Any]  					=							[
           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
       ]
       __lowercase   :   Any  					=							[
           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
       ]
       __lowercase   :   str  					=							[
           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    __A			( self					)						->	List[Any]:
           A_      =  JukeboxTokenizer.from_pretrained(self.model_id					)
           A_      =  tokenizer(**self.metas					)['''input_ids''']
           return tokens

       @slow
       def    __A			( self					)						->	List[str]:
           A_      =  JukeboxModel.from_pretrained(self.model_id       ,	min_duration=0					).eval()
           A_      =  self.prepare_inputs()

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

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

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

       @slow
       def    __A			( self					)						->	Optional[Any]:
           A_      =  False
           A_      =  JukeboxModel.from_pretrained(self.model_id       ,	min_duration=0					).eval()

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

           A_      =  model.priors[0]
           A_      =  0
           A_      =  top_prior.get_music_tokens_conds(_SCREAMING_SNAKE_CASE       ,	start=_SCREAMING_SNAKE_CASE       ,	end=start + top_prior.n_ctx					)
           A_      =  top_prior.get_metadata(labels[0].clone()       ,	_SCREAMING_SNAKE_CASE       ,	105_8304       ,	0					)

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

           A_  ,A_  ,A_      =  top_prior.get_cond(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)
           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					)						->	List[str]:
           A_      =  False

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

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

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

           A_      =  torch.cat((zs[1], torch.zeros(1       ,	2048 - zs[1].shape[-1]					))       ,	dim=-1					).long()
           A_      =  [upper_a, upper_a, model.vqvae.encode(_SCREAMING_SNAKE_CASE       ,	start_level=0       ,	bs_chunks=waveform.shape[0]					)[0]]
           A_      =  model._sample(
               _SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	sample_levels=[2]       ,	save_results=_SCREAMING_SNAKE_CASE       ,	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					)						->	Optional[int]:
           A_      =  JukeboxModel.from_pretrained(self.model_id       ,	min_duration=0					).eval()
           set_seed(0					)
           A_      =  torch.rand((1, 5120, 1)					)
           with torch.no_grad():
               A_      =  model.vqvae.encode(_SCREAMING_SNAKE_CASE       ,	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():
               A_      =  model.vqvae.decode(_SCREAMING_SNAKE_CASE       ,	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'''



       __lowercase   :   List[Any]  					=							(JukeboxModel,) if is_torch_available() else ()
       __lowercase   :   Optional[Any]  					=							'openai/jukebox-5b-lyrics'
       __lowercase   :   Any  					=							{
           '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
       __lowercase   :   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,
           1489, 1489, 1489, 1489, 1150, 1853, 1509, 1150, 1357, 1509, 6, 1272
       ]
       __lowercase   :   int  					=							[
           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
       ]

       __lowercase   :   Optional[Any]  					=							[
           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
       ]
       __lowercase   :   Tuple  					=							[
           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
       ]

       __lowercase   :   Dict  					=							[
           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
       ]
       __lowercase   :   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
       ]

       __lowercase   :   Optional[int]  					=							[
           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
       ]
       __lowercase   :   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, 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
       ]

       __lowercase   :   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
       ]
       __lowercase   :   str  					=							[
           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    __A			( self       ,	_SCREAMING_SNAKE_CASE					)						->	List[Any]:
           A_      =  JukeboxTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE					)
           A_      =  tokenizer(**self.metas					)['''input_ids''']
           return tokens

       @slow
       def    __A			( self					)						->	Optional[Any]:
           A_      =  JukeboxModel.from_pretrained(self.model_id       ,	min_duration=0					).eval()
           A_      =  self.prepare_inputs(self.model_id					)

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

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

           set_seed(0					)
           A_      =  model._sample(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	[2]       ,	sample_length=60 * model.priors[2].raw_to_tokens       ,	save_results=_SCREAMING_SNAKE_CASE					)
           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					)						->	Dict:
           A_      =  JukeboxModel.from_pretrained(self.model_id       ,	min_duration=0					).eval()
           A_      =  [i.cuda() for i in self.prepare_inputs(self.model_id					)]

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

       @slow
       def    __A			( self					)						->	Optional[Any]:
           A_      =  '''ArthurZ/jukebox_prior_0'''
           A_      =  JukeboxPrior.from_pretrained(_SCREAMING_SNAKE_CASE       ,	min_duration=0					).eval().half().to('''cuda'''					)

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







'''simple docstring'''




__snake_case  : str		       =    10**8





def 	_UpperCAmelCase       (				_UpperCamelCase       :      int  )      -> bool:
    if n % 10 == 0:
        return False
    A_      =  str(_UpperCamelCase  )
    return s == s[::-1]





def 	_UpperCAmelCase       (				)      -> int:
    A_      =  set()
    A_      =  1
    A_      =  5
    while sum_squares < LIMIT:
        A_      =  first_square + 1
        while sum_squares < LIMIT:
            if is_palindrome(_UpperCamelCase  ):
                answer.add(_UpperCamelCase  )
            last_square += 1
            sum_squares += last_square**2
        first_square += 1
        A_      =  first_square**2 + (first_square + 1) ** 2

    return sum(_UpperCamelCase  )


if __name__ == "__main__":
  print(solution())






