
'''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       lowerCamelCase_							(	unittest.TestCase  ):






   '''simple docstring'''

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

   A :     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
   ]

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

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

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

   A :     Optional[int]				=	[
       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
   ]
   A :     Dict				=	[
       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
   ]
   A :     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
   ]
   A :     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
   ]
   A :     Optional[Any]				=	[
       -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 :     List[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 :     List[str]				=	[
       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 :     Union[str, 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 	_UpperCamelCase       (  self ):
       UpperCamelCase__      =			JukeboxTokenizer.from_pretrained(self.model_id )
       UpperCamelCase__      =			tokenizer(**self.metas )["input_ids"]
       return tokens


   @slow
   def 	_UpperCamelCase       (  self ):
       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(__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 )
       UpperCamelCase__      =			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 )
       UpperCamelCase__      =			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 	_UpperCamelCase       (  self ):
       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(__UpperCamelCase     ,						start=__UpperCamelCase     ,						end=start + top_prior.n_ctx )
       UpperCamelCase__      =			top_prior.get_metadata(labels[0].clone()     ,						__UpperCamelCase     ,						1_05_83_04     ,						0 )

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

       UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__      =			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 	_UpperCamelCase       (  self ):
       UpperCamelCase__      =			False

       UpperCamelCase__      =			JukeboxModel.from_pretrained(self.model_id     ,						min_duration=0 ).eval()
       set_seed(0 )
       UpperCamelCase__      =			torch.rand((1, 51_20, 1) )
       UpperCamelCase__      =			list(self.prepare_inputs() )

       UpperCamelCase__      =			[model.vqvae.encode(__UpperCamelCase     ,						start_level=2     ,						bs_chunks=waveform.shape[0] )[0], None, None]
       UpperCamelCase__      =			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 ) )

       UpperCamelCase__      =			torch.cat((zs[0], torch.zeros(1     ,						20_48 - zs[0].shape[-1] ))     ,						dim=-1 ).long()
       UpperCamelCase__      =			[upper_a, model.vqvae.encode(__UpperCamelCase     ,						start_level=1     ,						bs_chunks=waveform.shape[0] )[0], None]
       UpperCamelCase__      =			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 ) )

       UpperCamelCase__      =			torch.cat((zs[1], torch.zeros(1     ,						20_48 - zs[1].shape[-1] ))     ,						dim=-1 ).long()
       UpperCamelCase__      =			[upper_a, upper_a, model.vqvae.encode(__UpperCamelCase     ,						start_level=0     ,						bs_chunks=waveform.shape[0] )[0]]
       UpperCamelCase__      =			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 	_UpperCamelCase       (  self ):
       UpperCamelCase__      =			JukeboxModel.from_pretrained(self.model_id     ,						min_duration=0 ).eval()
       set_seed(0 )
       UpperCamelCase__      =			torch.rand((1, 51_20, 1) )
       with torch.no_grad():
           UpperCamelCase__      =			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():
           UpperCamelCase__      =			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       lowerCamelCase_							(	unittest.TestCase  ):






   '''simple docstring'''

   A :     Any				=	(JukeboxModel,) if is_torch_available() else ()
   A :     Optional[Any]				=	"openai/jukebox-5b-lyrics"
   A :     str				=	{
       "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
   A :     List[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
   ]
   A :     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, 653, 653, 653, 653, 653, 653, 653, 653,
       653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
   ]

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

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

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

   A :     Optional[int]				=	[
       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
   ]
   A :     Tuple				=	[
       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 	_UpperCamelCase       (  self     ,						__UpperCamelCase ):
       UpperCamelCase__      =			JukeboxTokenizer.from_pretrained(__UpperCamelCase )
       UpperCamelCase__      =			tokenizer(**self.metas )["input_ids"]
       return tokens


   @slow
   def 	_UpperCamelCase       (  self ):
       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(__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 )
       UpperCamelCase__      =			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 )
       UpperCamelCase__      =			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 	_UpperCamelCase       (  self ):
       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(__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()
       UpperCamelCase__      =			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()
       UpperCamelCase__      =			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 	_UpperCamelCase       (  self ):
       UpperCamelCase__      =			"ArthurZ/jukebox_prior_0"
       UpperCamelCase__      =			JukeboxPrior.from_pretrained(__UpperCamelCase     ,						min_duration=0 ).eval().half().to("cuda" )

       UpperCamelCase__      =			self.prepare_inputs(__UpperCamelCase )[0].cuda()
       UpperCamelCase__      =			model.get_metadata(__UpperCamelCase     ,						0     ,						76_80     ,						0 )
       set_seed(0 )
       UpperCamelCase__      =			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'''







import heapq





def 						snake_case__ (lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_	):
    UpperCamelCase__      =			[(0, start)]  # cost from start node,end node
    UpperCamelCase__      =			set()
    while heap:
        ((UpperCamelCase__)    ,  (UpperCamelCase__))      =			heapq.heappop(lowerCAmelCase_	)
        if u in visited:
            continue
        visited.add(lowerCAmelCase_	)
        if u == end:
            return cost
        for v, c in graph[u]:
            if v in visited:
                continue
            UpperCamelCase__      =			cost + c
            heapq.heappush(lowerCAmelCase_					,  (next_item, v)	)
    return -1


lowercase__     :  Any    					=			{
    'A': [['B', 2], ['C', 5]],
    'B': [['A', 2], ['D', 3], ['E', 1], ['F', 1]],
    'C': [['A', 5], ['F', 3]],
    'D': [['B', 3]],
    'E': [['B', 4], ['F', 3]],
    'F': [['C', 3], ['E', 3]],
}
lowercase__     :  Optional[int]    					=			{
    'B': [['C', 1]],
    'C': [['D', 1]],
    'D': [['F', 1]],
    'E': [['B', 1], ['F', 3]],
    'F': [],
}
lowercase__     :  Union[str, Any]    					=			{
    'B': [['C', 1]],
    'C': [['D', 1]],
    'D': [['F', 1]],
    'E': [['B', 1], ['G', 2]],
    'F': [],
    'G': [['F', 1]],
}

lowercase__     :  Union[str, Any]    					=			dijkstra(G, 'E', 'C')
print(short_distance)  # E -- 3 --> F -- 3 --> C == 6

lowercase__     :  Optional[int]    					=			dijkstra(Ga, 'E', 'F')
print(short_distance)  # E -- 3 --> F == 3

lowercase__     :  Any    					=			dijkstra(Ga, 'E', 'F')
print(short_distance)  # E -- 2 --> G -- 1 --> F == 3

if __name__ == "__main__":
      import doctest

      doctest.testmod()




