



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     __snake_case						(					unittest.TestCase						):
							_a      				=			(JukeboxModel,) if is_torch_available() else ()
							_a      				=			'''openai/jukebox-1b-lyrics'''
							_a      				=			{
							    '''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      				=			[
							    1_864, 1_536, 1_213, 1_870, 1_357, 1_536, 519, 880, 1_323, 789, 1_082, 534,
							    1_000, 1_445, 1_105, 1_130, 967, 515, 1_434, 1_620, 534, 1_495, 283, 1_445,
							    333, 1_307, 539, 1_631, 1_528, 375, 1_434, 673, 627, 710, 778, 1_883,
							    1_405, 1_276, 1_455, 1_228
							]

							_a      				=			[
							    1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    653, 653, 653, 653, 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      				=			[
							    1_125, 1_751, 697, 1_776, 1_141, 1_476, 391, 697, 1_125, 684, 867, 416,
							    844, 1_372, 1_274, 717, 1_274, 844, 1_299, 1_419, 697, 1_370, 317, 1_125,
							    191, 1_440, 1_370, 1_440, 1_370, 282, 1_621, 1_370, 368, 349, 867, 1_872,
							    1_262, 869, 1_728, 747
							]
							_a      				=			[
							    416, 416, 1_125, 1_125, 416, 416, 416, 416, 416, 416, 416, 416,
							    416, 416, 416, 416, 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      				=			[
							    1_755, 842, 307, 1_843, 1_022, 1_395, 234, 1_554, 806, 739, 1_022, 442,
							    616, 556, 268, 1_499, 933, 457, 1_440, 1_837, 755, 985, 308, 902,
							    293, 1_443, 1_671, 1_141, 1_533, 555, 1_562, 1_061, 287, 417, 1_022, 2_008,
							    1_186, 1_015, 1_777, 268
							]
							_a      				=			[
							    854, 842, 1_353, 114, 1_353, 842, 185, 842, 185, 114, 591, 842,
							    185, 417, 185, 842, 307, 842, 591, 842, 185, 842, 307, 842,
							    591, 842, 1_353, 842, 185, 842, 591, 842, 591, 114, 591, 842,
							    185, 842, 591, 89
							]

							_a      				=			[1_058_304, 0, 786_432, 7_169, 507, 76, 27, 40, 30, 76]

							_a      				=			[
							    390, 1_160, 1_002, 1_907, 1_788, 1_788, 1_788, 1_907, 1_002, 1_002, 1_854, 1_002,
							    1_002, 1_002, 1_002, 1_002, 1_002, 1_160, 1_160, 1_606, 596, 596, 1_160, 1_002,
							    1_516, 596, 1_002, 1_002, 1_002, 1_907, 1_788, 1_788, 1_788, 1_854, 1_788, 1_907,
							    1_907, 1_788, 596, 1_626
							]
							_a      				=			[
							    1_236, 1_668, 1_484, 1_920, 1_848, 1_409, 139, 864, 1_828, 1_272, 1_599, 824,
							    1_672, 139, 555, 1_484, 824, 1_920, 555, 596, 1_579, 1_599, 1_231, 1_599,
							    1_637, 1_407, 212, 824, 1_599, 116, 1_433, 824, 258, 1_599, 1_433, 1_895,
							    1_063, 1_433, 1_433, 1_599
							]
							_a      				=			[
							    1_684, 1_873, 1_119, 1_189, 395, 611, 1_901, 972, 890, 1_337, 1_392, 1_927,
							    96, 972, 672, 780, 1_119, 890, 158, 771, 1_073, 1_927, 353, 1_331,
							    1_269, 1_459, 1_333, 1_645, 812, 1_577, 1_337, 606, 353, 981, 1_466, 619,
							    197, 391, 302, 1_930
							]
							_a      				=			[
							    390, 1_160, 1_002, 1_907, 1_788, 1_788, 1_788, 1_907, 1_002, 1_002, 1_854, 1_002,
							    1_002, 1_002, 1_002, 1_002, 1_002, 1_160, 1_160, 1_606, 596, 596, 1_160, 1_002,
							    1_516, 596, 1_002, 1_002, 1_002, 1_907, 1_788, 1_788, 1_788, 1_854, 1_788, 1_907,
							    1_907, 1_788, 596, 1_626
							]
							_a      				=			[
							    -0.0_4_9_2, -0.0_5_2_4, -0.0_5_6_5, -0.0_6_4_0, -0.0_6_8_6, -0.0_6_8_4, -0.0_6_7_7, -0.0_6_6_4,
							    -0.0_6_0_5, -0.0_4_9_0, -0.0_3_3_0, -0.0_1_6_8, -0.0_0_8_3, -0.0_0_7_5, -0.0_0_5_1, 0.0_0_2_5,
							    0.0_1_3_6, 0.0_2_6_1, 0.0_3_8_6, 0.0_4_9_7, 0.0_5_8_0, 0.0_5_9_9, 0.0_5_8_3, 0.0_6_1_4,
							    0.0_7_4_0, 0.0_8_8_9, 0.1_0_2_3, 0.1_1_6_2, 0.1_2_1_1, 0.1_2_1_2, 0.1_2_5_1, 0.1_3_3_6,
							    0.1_5_0_2, 0.1_6_8_6, 0.1_8_8_3, 0.2_1_4_8, 0.2_3_6_3, 0.2_4_5_8, 0.2_5_0_7, 0.2_5_3_1
							]
							_a      				=			[
							    0.0_2_5_6, -0.0_5_4_4, 0.1_6_0_0, -0.0_0_3_2, 0.1_0_6_6, 0.0_8_2_5, -0.0_0_1_3, 0.3_4_4_0,
							    0.0_2_1_0, 0.0_4_1_2, -0.1_7_7_7, -0.0_8_9_2, -0.0_1_6_4, 0.0_2_8_5, -0.0_6_1_3, -0.0_6_1_7,
							    -0.0_1_3_7, -0.0_2_0_1, -0.0_1_7_5, 0.0_2_1_5, -0.0_6_2_7, 0.0_5_2_0, -0.0_7_3_0, 0.0_9_7_0,
							    -0.0_1_0_0, 0.0_4_4_2, -0.0_5_8_6, 0.0_2_0_7, -0.0_0_1_5, -0.0_0_8_2
							]
							_a      				=			[
							    0.0_4_1_5, 0.0_8_7_7, 0.0_0_2_2, -0.0_0_5_5, 0.0_7_5_1, 0.0_3_3_4, 0.0_3_2_4, -0.0_0_6_8,
							    0.0_0_1_1, 0.0_0_1_7, -0.0_6_7_6, 0.0_6_5_5, -0.0_1_4_3, 0.0_3_9_9, 0.0_3_0_3, 0.0_7_4_3,
							    -0.0_1_6_8, -0.0_3_9_4, -0.1_1_1_3, 0.0_1_2_4, 0.0_4_4_2, 0.0_2_6_7, -0.0_0_0_3, -0.1_5_3_6,
							    -0.0_1_1_6, -0.1_8_3_7, -0.0_1_8_0, -0.1_0_2_6, -0.0_7_7_7, -0.0_4_5_6
							]
							_a      				=			[
							    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					:	Tuple):
									lowerCAmelCase_      :					Any     =			JukeboxTokenizer.from_pretrained(self.model_id)
									lowerCAmelCase_      :					Union[str, Any]     =			tokenizer(**self.metas)['''input_ids''']
									return tokens




							@slow
							def        UpperCAmelCase__      (	self					:	Any):
									lowerCAmelCase_      :					Tuple     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()
									lowerCAmelCase_      :					Tuple     =			self.prepare_inputs()

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

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

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




							@slow
							def        UpperCAmelCase__      (	self					:	List[Any]):
									lowerCAmelCase_      :					int     =			False
									lowerCAmelCase_      :					Any     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()

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

									lowerCAmelCase_      :					List[str]     =			model.priors[0]
									lowerCAmelCase_      :					str     =			0
									lowerCAmelCase_      :					Dict     =			top_prior.get_music_tokens_conds(A_    ,    start=A_    ,    end=start + top_prior.n_ctx)
									lowerCAmelCase_      :					Dict     =			top_prior.get_metadata(labels[0].clone()    ,    A_    ,    1_0_5_8_3_0_4    ,    0)

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

									lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_      :					List[str]     =			top_prior.get_cond(A_    ,    A_)
									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        UpperCAmelCase__      (	self					:	Union[str, Any]):
									lowerCAmelCase_      :					int     =			False

									lowerCAmelCase_      :					Optional[Any]     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()
									set_seed(0)
									lowerCAmelCase_      :					Dict     =			torch.rand((1, 5_1_2_0, 1))
									lowerCAmelCase_      :					List[str]     =			list(self.prepare_inputs())

									lowerCAmelCase_      :					List[str]     =			[model.vqvae.encode(A_    ,    start_level=2    ,    bs_chunks=waveform.shape[0])[0], None, None]
									lowerCAmelCase_      :					Tuple     =			model._sample(
									    A_    ,    A_    ,    sample_levels=[0]    ,    save_results=A_    ,    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))

									lowerCAmelCase_      :					Dict     =			torch.cat((zs[0], torch.zeros(1    ,    2_0_4_8 - zs[0].shape[-1]))    ,    dim=-1).long()
									lowerCAmelCase_      :					Optional[int]     =			[upper_a, model.vqvae.encode(A_    ,    start_level=1    ,    bs_chunks=waveform.shape[0])[0], None]
									lowerCAmelCase_      :					List[Any]     =			model._sample(
									    A_    ,    A_    ,    sample_levels=[1]    ,    save_results=A_    ,    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))

									lowerCAmelCase_      :					Dict     =			torch.cat((zs[1], torch.zeros(1    ,    2_0_4_8 - zs[1].shape[-1]))    ,    dim=-1).long()
									lowerCAmelCase_      :					Tuple     =			[upper_a, upper_a, model.vqvae.encode(A_    ,    start_level=0    ,    bs_chunks=waveform.shape[0])[0]]
									lowerCAmelCase_      :					Optional[Any]     =			model._sample(
									    A_    ,    A_    ,    sample_levels=[2]    ,    save_results=A_    ,    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        UpperCAmelCase__      (	self					:	List[Any]):
									lowerCAmelCase_      :					str     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()
									set_seed(0)
									lowerCAmelCase_      :					List[Any]     =			torch.rand((1, 5_1_2_0, 1))
									with torch.no_grad():
											lowerCAmelCase_      :					Dict     =			model.vqvae.encode(A_    ,    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():
											lowerCAmelCase_      :					Dict     =			model.vqvae.decode(A_    ,    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     __snake_case						(					unittest.TestCase						):
							_a      				=			(JukeboxModel,) if is_torch_available() else ()
							_a      				=			'''openai/jukebox-5b-lyrics'''
							_a      				=			{
							    '''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      				=			[
							    1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    1_489, 1_489, 1_489, 1_489, 1_150, 1_853, 1_509, 1_150, 1_357, 1_509, 6, 1_272
							]
							_a      				=			[
							    1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    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      				=			[
							    1_125, 416, 1_125, 1_125, 1_125, 1_125, 1_125, 416, 416, 416, 416, 416,
							    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
							    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
							    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      				=			[
							    416, 416, 1_125, 1_125, 416, 416, 416, 416, 416, 416, 416, 416, 416,
							    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
							    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 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      				=			[
							    1_755, 1_061, 234, 1_755, 1_061, 1_755, 185, 290, 307, 307, 616, 616,
							    616, 616, 616, 616, 307, 290, 417, 1_755, 234, 1_755, 185, 290,
							    290, 290, 307, 616, 616, 616, 616, 616, 290, 234, 234, 1_755,
							    234, 234, 1_755, 234, 185, 185, 307, 616, 616, 616, 616, 290,
							    1_755, 1_755, 1_755, 234, 234, 1_755, 1_572, 290, 307, 616, 34, 616
							]
							_a      				=			[
							    854, 842, 1_353, 114, 1_353, 842, 185, 842, 185, 114, 591, 842, 185,
							    417, 185, 842, 307, 842, 591, 842, 185, 842, 185, 842, 591, 842,
							    1_353, 842, 185, 842, 591, 842, 591, 114, 591, 842, 185, 842, 591,
							    89, 591, 842, 591, 842, 591, 417, 1_372, 842, 1_372, 842, 34, 842,
							    185, 89, 591, 842, 185, 842, 591, 632
							]

							_a      				=			[
							    1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    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      				=			[
							    1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
							    653, 653, 653, 653, 653, 653, 653, 1_853, 1_177, 1_536, 1_228,
							    710, 475, 1_489, 1_229, 1_224, 231, 1_224, 252, 1_434, 653, 475,
							    1_106, 1_877, 1_599, 1_228, 1_600, 1_683, 1_182, 1_853, 475, 1_864,
							    252, 1_229, 1_434, 2_001
							]

							_a      				=			[
							    1_125, 1_125, 416, 1_125, 1_125, 416, 1_125, 1_125, 416, 416, 1_125, 416,
							    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
							    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
							    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      				=			[
							    491, 1_755, 34, 1_613, 1_755, 417, 992, 1_613, 222, 842, 1_353, 1_613,
							    844, 632, 185, 1_613, 844, 632, 185, 1_613, 185, 842, 677, 1_613,
							    185, 114, 1_353, 1_613, 307, 89, 844, 1_613, 307, 1_332, 234, 1_979,
							    307, 89, 1_353, 616, 34, 842, 185, 842, 34, 842, 185, 842,
							    307, 114, 185, 89, 34, 1_268, 185, 89, 34, 842, 185, 89
							]




							def        UpperCAmelCase__      (	self					:	int    ,    A_					:	Union[str, Any]):
									lowerCAmelCase_      :					Tuple     =			JukeboxTokenizer.from_pretrained(A_)
									lowerCAmelCase_      :					Any     =			tokenizer(**self.metas)['''input_ids''']
									return tokens




							@slow
							def        UpperCAmelCase__      (	self					:	Dict):
									lowerCAmelCase_      :					Optional[Any]     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()
									lowerCAmelCase_      :					Union[str, Any]     =			self.prepare_inputs(self.model_id)

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

									set_seed(0)
									lowerCAmelCase_      :					Optional[int]     =			model._sample(A_    ,    A_    ,    [1]    ,    sample_length=6_0 * model.priors[1].raw_to_tokens    ,    save_results=A_)
									self.assertIn(zs[1][0].detach().cpu().tolist()    ,    [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2])

									set_seed(0)
									lowerCAmelCase_      :					Tuple     =			model._sample(A_    ,    A_    ,    [2]    ,    sample_length=6_0 * model.priors[2].raw_to_tokens    ,    save_results=A_)
									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					:	List[str]):
									lowerCAmelCase_      :					str     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()
									lowerCAmelCase_      :					Optional[Any]     =			[i.cuda() for i in self.prepare_inputs(self.model_id)]

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







							@slow
							def        UpperCAmelCase__      (	self					:	Optional[int]):
									lowerCAmelCase_      :					Tuple     =			'''ArthurZ/jukebox_prior_0'''
									lowerCAmelCase_      :					Any     =			JukeboxPrior.from_pretrained(A_    ,    min_duration=0).eval().half().to('''cuda''')

									lowerCAmelCase_      :					Any     =			self.prepare_inputs(A_)[0].cuda()
									lowerCAmelCase_      :					Any     =			model.get_metadata(A_    ,    0    ,    7_6_8_0    ,    0)
									set_seed(0)
									lowerCAmelCase_      :					Optional[int]     =			model.sample(1    ,    metadata=A_    ,    sample_tokens=6_0)
									self.assertIn(outputs[0].cpu().tolist()    ,    [self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2])










from math import radians as angle_to_radians
from math import sin

# Acceleration Constant on Earth (unit m/s^2)
A__  :  Any       = 9.8_0_6_6_5


def       UpperCamelCase(  __UpperCamelCase							:		float					,__UpperCamelCase							:		float			):

		# Ensure valid instance
		if not isinstance(__UpperCamelCase					,(int, float)			):
				raise TypeError('''Invalid velocity. Should be a positive number.'''			)

		if not isinstance(__UpperCamelCase					,(int, float)			):
				raise TypeError('''Invalid angle. Range is 1-90 degrees.'''			)

		# Ensure valid angle
		if angle > 90 or angle < 1:
				raise ValueError('''Invalid angle. Range is 1-90 degrees.'''			)

		# Ensure valid velocity
		if init_velocity < 0:
				raise ValueError('''Invalid velocity. Should be a positive number.'''			)


def       UpperCamelCase(  __UpperCamelCase							:		float					,__UpperCamelCase							:		float			):
		check_args(__UpperCamelCase					,__UpperCamelCase			)
		lowerCAmelCase_      :					Optional[Any]     =			angle_to_radians(2 * angle			)
		return round(init_velocity**2 * sin(__UpperCamelCase			) / g					,2			)


def       UpperCamelCase(  __UpperCamelCase							:		float					,__UpperCamelCase							:		float			):
		check_args(__UpperCamelCase					,__UpperCamelCase			)
		lowerCAmelCase_      :					Any     =			angle_to_radians(__UpperCamelCase			)
		return round(init_velocity**2 * sin(__UpperCamelCase			) ** 2 / (2 * g)					,2			)


def       UpperCamelCase(  __UpperCamelCase							:		float					,__UpperCamelCase							:		float			):
		check_args(__UpperCamelCase					,__UpperCamelCase			)
		lowerCAmelCase_      :					Dict     =			angle_to_radians(__UpperCamelCase			)
		return round(2 * init_velocity * sin(__UpperCamelCase			) / g					,2			)


def       UpperCamelCase(  ):
		lowerCAmelCase_ ,  lowerCAmelCase_      :					Any     =			25, 20
		assert horizontal_distance(__UpperCamelCase					,__UpperCamelCase			) == 4_0.9_7
		assert max_height(__UpperCamelCase					,__UpperCamelCase			) == 3.7_3
		assert total_time(__UpperCamelCase					,__UpperCamelCase			) == 1.7_4


if __name__ == "__main__":
						from doctest import testmod

						testmod()

						# Get input from user
						A__  :  List[Any]       = float(input('''Initial Velocity: ''').strip())

						# Get input from user
						A__  :  Dict       = float(input('''angle: ''').strip())

						# Print results
						print()
						print('''Results: ''')
						print(F'''Horizontal Distance: {horizontal_distance(init_vel, angle)!s} [m]''')
						print(F'''Maximum Height: {max_height(init_vel, angle)!s} [m]''')
						print(F'''Total Time: {total_time(init_vel, angle)!s} [s]''')





