code: 
  
  
  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 		_SCREAMING_SNAKE_CASE    (	unittest.TestCase    ):
  
  
       @slow
       @unittest.skip("""Model is not available."""						)
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       Optional[Any]       =    """nvidia/megatron-gpt2-345m/"""
         if "MYDIR" in os.environ:
           _UpperCAmelCase		:       List[str]       =    os.path.join(os.environ["""MYDIR"""]       ,				A_						)
         _UpperCAmelCase		:       str       =    GPTaLMHeadModel.from_pretrained(A_						)
         model.to(A_						)
         model.half()
  
         _UpperCAmelCase		:       Dict       =    torch.tensor(
             [[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]       ,				device=A_       ,				dtype=torch.long       ,				)
  
         with torch.no_grad():
           _UpperCAmelCase		:       Optional[int]       =    model(A_						).logits
  
         _UpperCAmelCase		:       List[Any]       =    torch.Size((1, 9, 5_02_57)						)
         self.assertEqual(output.shape       ,				A_						)
  
         _UpperCAmelCase		:       Optional[Any]       =    torch.tensor(
             [
                 4.9_4_1_4,
                 -0.2_9_2_0,
                 -1.2_1_4_8,
                 -4.0_2_7_3,
                 -0.5_1_6_1,
                 -5.2_1_0_9,
                 -1.2_4_1_2,
                 -1.8_3_0_1,
                 -1.7_7_3_4,
                 -4.7_1_4_8,
                 -0.2_3_1_7,
                 -1.0_8_1_1,
                 -2.1_7_7_7,
                 0.4_1_4_1,
                 -3.7_9_6_9,
                 -4.0_5_8_6,
                 -2.5_3_3_2,
                 -3.3_8_0_9,
                 4.3_8_6_7,
             ]       ,				device=A_       ,				dtype=torch.half       ,				)
  
         for i in range(19						):
           _UpperCAmelCase,_UpperCAmelCase		:       int       =    8 * i // 17, 27_92 * i  # along the diagonal
           _UpperCAmelCase,_UpperCAmelCase		:       List[Any]       =    output[0, r, c], expected_diag[i]
           _UpperCAmelCase		:       Tuple       =    F'''row={r} col={c} computed={computed} expected={expected}'''
           self.assertAlmostEqual(A_       ,				A_       ,				delta=1e-4       ,				msg=A_						)
  
  
  
  
  
  
style_context: 
  
  
  from typing import Any
  
  
  class 		_SCREAMING_SNAKE_CASE    :
  
  
       def __init__(     self						):
         _UpperCAmelCase		:       Union[str, Any]       =    []
         _UpperCAmelCase		:       Dict       =    0
  
  
       def __str__(     self						):
         _UpperCAmelCase		:       Union[str, Any]       =    """<""" + str(self.stack						)[1:-1] + """>"""
         return printed
  
  
       def 					__snake_case(     self       ,				A_						):
         self.stack.append(A_						)
         _UpperCAmelCase		:       Optional[Any]       =    self.length + 1
  
  
       def 					__snake_case(     self						):
         self.rotate(1						)
         _UpperCAmelCase		:       List[Any]       =    self.stack[self.length - 1]
         _UpperCAmelCase		:       Optional[Any]       =    self.stack[:-1]
         self.rotate(self.length - 1						)
         _UpperCAmelCase		:       Union[str, Any]       =    self.length - 1
         return dequeued
  
  
       def 					__snake_case(     self       ,				A_						):
         for _ in range(A_						):
           _UpperCAmelCase		:       Dict       =    self.stack[0]
           _UpperCAmelCase		:       List[Any]       =    self.stack[1:]
           self.put(A_						)
           _UpperCAmelCase		:       Dict       =    self.length - 1
  
  
       def 					__snake_case(     self						):
         _UpperCAmelCase		:       Tuple       =    self.get()
         self.put(A_						)
         self.rotate(self.length - 1						)
         return front
  
  
       def 					__snake_case(     self						):
         return self.length
  
  
  
  
  
  