





"""simple docstring"""




import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
						import torch

						from transformers import GPTaLMHeadModel




@require_torch
@require_sentencepiece
@require_tokenizers
class 				_a		(   unittest.TestCase       ):






				@slow
				@unittest.skip('''Model is not available.'''							)
				def        snake_case			(				self	:					Tuple							)			->   List[Any]:




										'''simple docstring'''



										_UpperCamelCase    :      Dict						     =					'''nvidia/megatron-gpt2-345m/'''
										if "MYDIR" in os.environ:
																_UpperCamelCase    :      Optional[Any]						     =					os.path.join(os.environ['''MYDIR'''],				lowerCAmelCase__							)
										_UpperCamelCase    :      Any						     =					GPTaLMHeadModel.from_pretrained(lowerCAmelCase__							)
										model.to(lowerCAmelCase__							)
										model.half()

										_UpperCamelCase    :      str						     =					torch.tensor(
										    [[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]],				device=lowerCAmelCase__,				dtype=torch.long,				)

										with torch.no_grad():
																_UpperCamelCase    :      int						     =					model(lowerCAmelCase__							).logits

										_UpperCamelCase    :      Optional[int]						     =					torch.Size((1, 9, 5_0_2_5_7)							)
										self.assertEqual(output.shape,				lowerCAmelCase__							)

										_UpperCamelCase    :      Dict						     =					torch.tensor(
										    [
										        4.9_414,
										        -0.2_920,
										        -1.2_148,
										        -4.0_273,
										        -0.5_161,
										        -5.2_109,
										        -1.2_412,
										        -1.8_301,
										        -1.7_734,
										        -4.7_148,
										        -0.2_317,
										        -1.0_811,
										        -2.1_777,
										        0.4_141,
										        -3.7_969,
										        -4.0_586,
										        -2.5_332,
										        -3.3_809,
										        4.3_867,
										    ],				device=lowerCAmelCase__,				dtype=torch.half,				)

										for i in range(1_9							):
																_UpperCamelCase    :      str						     =					8 * i // 1_7, 2_7_9_2 * i  # along the diagonal
																_UpperCamelCase    :      Optional[int]						     =					output[0, r, c], expected_diag[i]
																_UpperCamelCase    :      Any						     =					f"""row={r} col={c} computed={computed} expected={expected}"""
																self.assertAlmostEqual(lowerCAmelCase__,				lowerCAmelCase__,				delta=1e-4,				msg=lowerCAmelCase__							)










"""simple docstring"""




def  a_       (							_lowercase     ):
						if len(_lowercase     ) > 1:
												_UpperCamelCase    :      List[str]						     =					len(_lowercase     ) // 2  # Finds the middle of the array
												_UpperCamelCase    :      Dict						     =					arr[
												    :middle_length
												]  # Creates an array of the elements in the first half.
												_UpperCamelCase    :      Union[str, Any]						     =					arr[
												    middle_length:
												]  # Creates an array of the elements in the second half.
												_UpperCamelCase    :      str						     =					len(_lowercase     )
												_UpperCamelCase    :      List[str]						     =					len(_lowercase     )
												merge(_lowercase     )  # Starts sorting the left.
												merge(_lowercase     )  # Starts sorting the right
												_UpperCamelCase    :      Dict						     =					0  # Left Counter
												_UpperCamelCase    :      Union[str, Any]						     =					0  # Right Counter
												_UpperCamelCase    :      Optional[int]						     =					0  # Position Counter
												while (
												    left_index < left_size and right_index < right_size
												):  # Runs until the lowers size of the left and right are sorted.
																		if left_array[left_index] < right_array[right_index]:
																								_UpperCamelCase    :      Dict						     =					left_array[left_index]
																								left_index += 1
																		else:
																								_UpperCamelCase    :      Tuple						     =					right_array[right_index]
																								right_index += 1
																		index += 1
												while (
												    left_index < left_size
												):  # Adds the left over elements in the left half of the array
																		_UpperCamelCase    :      int						     =					left_array[left_index]
																		left_index += 1
																		index += 1
												while (
												    right_index < right_size
												):  # Adds the left over elements in the right half of the array
																		_UpperCamelCase    :      Any						     =					right_array[right_index]
																		right_index += 1
																		index += 1
						return arr


if __name__ == "__main__":
						import doctest

						doctest.testmod()



