



import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
    import jax.numpy as jnp

    from transformers.models.albert.modeling_flax_albert import (
        FlaxAlbertForMaskedLM,
        FlaxAlbertForMultipleChoice,
        FlaxAlbertForPreTraining,
        FlaxAlbertForQuestionAnswering,
        FlaxAlbertForSequenceClassification,
        FlaxAlbertForTokenClassification,
        FlaxAlbertModel,
    )





class 		UpperCAmelCase						(						unittest.TestCase				):





   def __init__(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase=13       ,     _lowerCAmelCase=7       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=99       ,     _lowerCAmelCase=32       ,     _lowerCAmelCase=5       ,     _lowerCAmelCase=4       ,     _lowerCAmelCase=37       ,     _lowerCAmelCase="gelu"       ,     _lowerCAmelCase=0.1       ,     _lowerCAmelCase=0.1       ,     _lowerCAmelCase=512       ,     _lowerCAmelCase=16       ,     _lowerCAmelCase=2       ,     _lowerCAmelCase=0.02       ,     _lowerCAmelCase=4       ,     ):
         _lowerCAmelCase 			=  parent
         _lowerCAmelCase 			=  batch_size
         _lowerCAmelCase 			=  seq_length
         _lowerCAmelCase 			=  is_training
         _lowerCAmelCase 			=  use_attention_mask
         _lowerCAmelCase 			=  use_token_type_ids
         _lowerCAmelCase 			=  use_labels
         _lowerCAmelCase 			=  vocab_size
         _lowerCAmelCase 			=  hidden_size
         _lowerCAmelCase 			=  num_hidden_layers
         _lowerCAmelCase 			=  num_attention_heads
         _lowerCAmelCase 			=  intermediate_size
         _lowerCAmelCase 			=  hidden_act
         _lowerCAmelCase 			=  hidden_dropout_prob
         _lowerCAmelCase 			=  attention_probs_dropout_prob
         _lowerCAmelCase 			=  max_position_embeddings
         _lowerCAmelCase 			=  type_vocab_size
         _lowerCAmelCase 			=  type_sequence_label_size
         _lowerCAmelCase 			=  initializer_range
         _lowerCAmelCase 			=  num_choices





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size			)

         _lowerCAmelCase 			=  None
         if self.use_attention_mask:
               _lowerCAmelCase 			=  random_attention_mask([self.batch_size, self.seq_length]			)

         _lowerCAmelCase 			=  None
         if self.use_token_type_ids:
               _lowerCAmelCase 			=  ids_tensor([self.batch_size, self.seq_length]       ,     self.type_vocab_size			)

         _lowerCAmelCase 			=  AlbertConfig(
             vocab_size=self.vocab_size       ,     hidden_size=self.hidden_size       ,     num_hidden_layers=self.num_hidden_layers       ,     num_attention_heads=self.num_attention_heads       ,     intermediate_size=self.intermediate_size       ,     hidden_act=self.hidden_act       ,     hidden_dropout_prob=self.hidden_dropout_prob       ,     attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,     max_position_embeddings=self.max_position_embeddings       ,     type_vocab_size=self.type_vocab_size       ,     is_decoder=_lowerCAmelCase       ,     initializer_range=self.initializer_range       ,     )

         return config, input_ids, token_type_ids, attention_mask





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.prepare_config_and_inputs()
         _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase 			=  config_and_inputs
         _lowerCAmelCase 			=  {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
         return config, inputs_dict






@require_flax
class 		UpperCAmelCase						(						snake_case_       ,unittest.TestCase				):
   SCREAMING_SNAKE_CASE__							=     (
       (
           FlaxAlbertModel,
           FlaxAlbertForPreTraining,
           FlaxAlbertForMaskedLM,
           FlaxAlbertForMultipleChoice,
           FlaxAlbertForQuestionAnswering,
           FlaxAlbertForSequenceClassification,
           FlaxAlbertForTokenClassification,
           FlaxAlbertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  FlaxAlbertModelTester(self			)





   @slow
   def 	__lowerCAmelCase							(	self			):
         for model_class_name in self.all_model_classes:
               _lowerCAmelCase 			=  model_class_name.from_pretrained('''albert-base-v2'''			)
               _lowerCAmelCase 			=  model(np.ones((1, 1)			)			)
               self.assertIsNotNone(_lowerCAmelCase			)






@require_flax
class 		UpperCAmelCase						(						unittest.TestCase				):





   @slow
   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  FlaxAlbertModel.from_pretrained('''albert-base-v2'''			)
         _lowerCAmelCase 			=  np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]			)
         _lowerCAmelCase 			=  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]			)
         _lowerCAmelCase 			=  model(_lowerCAmelCase       ,     attention_mask=_lowerCAmelCase			)[0]
         _lowerCAmelCase 			=  (1, 11, 768)
         self.assertEqual(output.shape       ,     _lowerCAmelCase			)
         _lowerCAmelCase 			=  np.array(
             [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]]			)

         self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]       ,     _lowerCAmelCase       ,     atol=1E-4			)			)











import functools
import gc
import inspect

import torch

from .imports import is_npu_available, is_xpu_available
def 							UpperCAmelCase__		(				*_SCREAMING_SNAKE_CASE   :	Tuple				)->List[Any]:
      if not isinstance(_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE				):
            _lowerCAmelCase 			=  list(_SCREAMING_SNAKE_CASE				)
      for i in range(len(_SCREAMING_SNAKE_CASE				)				):
            _lowerCAmelCase 			=  None
      gc.collect()
      if is_xpu_available():
            torch.xpu.empty_cache()
      elif is_npu_available():
            torch.npu.empty_cache()
      else:
            torch.cuda.empty_cache()
      return objects
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	Exception				)->bool:
      _lowerCAmelCase 			=  [
          '''CUDA out of memory.''',  # CUDA OOM
          '''cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.''',  # CUDNN SNAFU
          '''DefaultCPUAllocator: can\'t allocate memory''',  # CPU OOM
      ]
      if isinstance(_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE				) and len(exception.args				) == 1:
            return any(err in exception.args[0] for err in _statements				)
      return False


def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	callable = None     ,				_SCREAMING_SNAKE_CASE   :	int = 1_2_8				)->Optional[int]:
      if function is None:
            return functools.partial(_SCREAMING_SNAKE_CASE     ,				starting_batch_size=_SCREAMING_SNAKE_CASE				)

      _lowerCAmelCase 			=  starting_batch_size

      def decorator(*_SCREAMING_SNAKE_CASE   :	Optional[int]     ,				**_SCREAMING_SNAKE_CASE   :	Optional[Any]				):
            nonlocal batch_size
            gc.collect()
            if is_xpu_available():
                  torch.xpu.empty_cache()
            elif is_npu_available():
                  torch.npu.empty_cache()
            else:
                  torch.cuda.empty_cache()
            _lowerCAmelCase 			=  list(inspect.signature(_SCREAMING_SNAKE_CASE				).parameters.keys()				)
            # Guard against user error
            if len(_SCREAMING_SNAKE_CASE				) < (len(_SCREAMING_SNAKE_CASE				) + 1):
                  _lowerCAmelCase 			=  ''', '''.join([f'''{arg}={value}''' for arg, value in zip(params[1:]     ,				args[1:]				)]				)
                  raise TypeError(
                      f'''Batch size was passed into `{function.__name__}` as the first argument when called.'''
                      f'''Remove this as the decorator already does so: `{function.__name__}({arg_str})`'''				)
            while True:
                  if batch_size == 0:
                        raise RuntimeError('''No executable batch size found, reached zero.'''				)
                  try:
                        return function(_SCREAMING_SNAKE_CASE     ,				*_SCREAMING_SNAKE_CASE     ,				**_SCREAMING_SNAKE_CASE				)
                  except Exception as e:
                        if should_reduce_batch_size(_SCREAMING_SNAKE_CASE				):
                              gc.collect()
                              if is_xpu_available():
                                    torch.xpu.empty_cache()
                              elif is_npu_available():
                                    torch.npu.empty_cache()
                              else:
                                    torch.cuda.empty_cache()
                              batch_size //= 2
                        else:
                              raise

      return decorator