





import logging

# standard libraries
import sys

import numpy as np
import pytest  # type: ignore

# Custom/local libraries
from matrix import matrix_operation as matop

__lowercase				    =      [[12, 10], [3, 9]]
__lowercase				    =      [[3, 4], [7, 4]]
__lowercase				    =      [[3, 0, 2], [2, 0, -2], [0, 1, 1]]
__lowercase				    =      [[3, 0, -2], [2, 0, 2], [0, 1, 1]]
__lowercase				    =      [[3, 0, 2], [2, 0, -2], [0, 1, 1], [2, 0, -2]]
__lowercase				    =      [1]
__lowercase				    =      [2]

__lowercase				    =      logging.getLogger()
__lowercase				    =      logging.DEBUG
__lowercase				    =      logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)


@pytest.mark.mat_ops
@pytest.mark.parametrize(
    ('''mat1''', '''mat2''')  ,				[(mat_a, mat_b), (mat_c, mat_d), (mat_d, mat_e), (mat_f, mat_h)])
def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Tuple  ,				__lowercase:		Optional[Any]):



  '''simple docstring'''






  if (np.array(__lowercase)).shape < (2, 2) or (np.array(__lowercase)).shape < (2, 2):
    with pytest.raises(__lowercase):
      logger.info(F'''\n\t{test_addition.__name__} returned integer''')
      matop.add(__lowercase  ,				__lowercase)
  elif (np.array(__lowercase)).shape == (np.array(__lowercase)).shape:
    logger.info(F'''\n\t{test_addition.__name__} with same matrix dims''')
    __lowerCAmelCase								=					(np.array(__lowercase) + np.array(__lowercase)).tolist()
    __lowerCAmelCase								=					matop.add(__lowercase  ,				__lowercase)
    assert theo == act
  else:
    with pytest.raises(__lowercase):
      logger.info(F'''\n\t{test_addition.__name__} with different matrix dims''')
      matop.add(__lowercase  ,				__lowercase)


@pytest.mark.mat_ops
@pytest.mark.parametrize(
    ('''mat1''', '''mat2''')  ,				[(mat_a, mat_b), (mat_c, mat_d), (mat_d, mat_e), (mat_f, mat_h)])
def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Tuple  ,				__lowercase:		Optional[Any]):



  '''simple docstring'''






  if (np.array(__lowercase)).shape < (2, 2) or (np.array(__lowercase)).shape < (2, 2):
    with pytest.raises(__lowercase):
      logger.info(F'''\n\t{test_subtraction.__name__} returned integer''')
      matop.subtract(__lowercase  ,				__lowercase)
  elif (np.array(__lowercase)).shape == (np.array(__lowercase)).shape:
    logger.info(F'''\n\t{test_subtraction.__name__} with same matrix dims''')
    __lowerCAmelCase								=					(np.array(__lowercase) - np.array(__lowercase)).tolist()
    __lowerCAmelCase								=					matop.subtract(__lowercase  ,				__lowercase)
    assert theo == act
  else:
    with pytest.raises(__lowercase):
      logger.info(F'''\n\t{test_subtraction.__name__} with different matrix dims''')
      assert matop.subtract(__lowercase  ,				__lowercase)


@pytest.mark.mat_ops
@pytest.mark.parametrize(
    ('''mat1''', '''mat2''')  ,				[(mat_a, mat_b), (mat_c, mat_d), (mat_d, mat_e), (mat_f, mat_h)])
def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Optional[int]  ,				__lowercase:		Optional[int]):



  '''simple docstring'''






  if (np.array(__lowercase)).shape < (2, 2) or (np.array(__lowercase)).shape < (2, 2):
    logger.info(F'''\n\t{test_multiplication.__name__} returned integer''')
    with pytest.raises(__lowercase):
      matop.add(__lowercase  ,				__lowercase)
  elif (np.array(__lowercase)).shape == (np.array(__lowercase)).shape:
    logger.info(F'''\n\t{test_multiplication.__name__} meets dim requirements''')
    __lowerCAmelCase								=					(np.matmul(__lowercase  ,				__lowercase)).tolist()
    __lowerCAmelCase								=					matop.multiply(__lowercase  ,				__lowercase)
    assert theo == act
  else:
    with pytest.raises(__lowercase):
      logger.info(
          F'''\n\t{test_multiplication.__name__} does not meet dim requirements''')
      assert matop.subtract(__lowercase  ,				__lowercase)


@pytest.mark.mat_ops
def        __SCREAMING_SNAKE_CASE	(				):



  '''simple docstring'''






  __lowerCAmelCase								=					(3.5 * np.array(__lowercase)).tolist()
  __lowerCAmelCase								=					matop.scalar_multiply(__lowercase  ,				3.5)
  assert theo == act


@pytest.mark.mat_ops
def        __SCREAMING_SNAKE_CASE	(				):



  '''simple docstring'''






  __lowerCAmelCase								=					(np.identity(5)).tolist()
  __lowerCAmelCase								=					matop.identity(5)
  assert theo == act




@pytest.mark.mat_ops
@pytest.mark.parametrize('''mat'''  ,				[mat_a, mat_b, mat_c, mat_d, mat_e, mat_f])
def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Optional[Any]):



  '''simple docstring'''






  if (np.array(__lowercase)).shape < (2, 2):
    with pytest.raises(__lowercase):
      logger.info(F'''\n\t{test_transpose.__name__} returned integer''')
      matop.transpose(__lowercase)
  else:
    __lowerCAmelCase								=					(np.transpose(__lowercase)).tolist()
    __lowerCAmelCase								=					matop.transpose(__lowercase  ,				return_map=__lowercase)
    assert theo == act







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






    '''simple docstring'''





    @slow
    @unittest.skip('''Model is not available.'''				)
    def 						A		( self    :		str				):

      """simple docstring"""





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

      __lowerCAmelCase								=					torch.tensor(
          [[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]      ,    device=snake_case      ,    dtype=torch.long      ,    )

      with torch.no_grad():
        __lowerCAmelCase								=					model(snake_case				).logits

      __lowerCAmelCase								=					torch.Size((1, 9, 50_257)				)
      self.assertEqual(output.shape      ,    snake_case				)

      __lowerCAmelCase								=					torch.tensor(
          [
              4.9414,
              -0.2920,
              -1.2148,
              -4.0273,
              -0.5161,
              -5.2109,
              -1.2412,
              -1.8301,
              -1.7734,
              -4.7148,
              -0.2317,
              -1.0811,
              -2.1777,
              0.4141,
              -3.7969,
              -4.0586,
              -2.5332,
              -3.3809,
              4.3867,
          ]      ,    device=snake_case      ,    dtype=torch.half      ,    )

      for i in range(19				):
        __lowerCAmelCase						,    __lowerCAmelCase								=					8 * i // 17, 2_792 * i  # along the diagonal
        __lowerCAmelCase						,    __lowerCAmelCase								=					output[0, r, c], expected_diag[i]
        __lowerCAmelCase								=					f'''row={r} col={c} computed={computed} expected={expected}'''
        self.assertAlmostEqual(snake_case      ,    snake_case      ,    delta=1e-4      ,    msg=snake_case				)
