





'''simple docstring'''

import unittest

from transformers import BigBirdConfig, 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

   from transformers.models.big_bird.modeling_flax_big_bird import (
       FlaxBigBirdForCausalLM,
       FlaxBigBirdForMaskedLM,
       FlaxBigBirdForMultipleChoice,
       FlaxBigBirdForPreTraining,
       FlaxBigBirdForQuestionAnswering,
       FlaxBigBirdForSequenceClassification,
       FlaxBigBirdForTokenClassification,
       FlaxBigBirdModel,
   )
class 				__lowerCAmelCase     (    unittest.TestCase ):






       """simple docstring"""





       def __init__(							self   :		List[Any]       , lowerCAmelCase__   :		Optional[Any]       , lowerCAmelCase__   :		Union[str, Any]=2       , lowerCAmelCase__   :		str=56       , lowerCAmelCase__   :		Union[str, Any]=True       , lowerCAmelCase__   :		Optional[Any]=True       , lowerCAmelCase__   :		Optional[int]=True       , lowerCAmelCase__   :		Tuple=True       , lowerCAmelCase__   :		Union[str, Any]=99       , lowerCAmelCase__   :		Optional[Any]=32       , lowerCAmelCase__   :		List[Any]=2       , lowerCAmelCase__   :		Union[str, Any]=2       , lowerCAmelCase__   :		Optional[int]=7       , lowerCAmelCase__   :		str="gelu_new"       , lowerCAmelCase__   :		List[str]=0.1       , lowerCAmelCase__   :		Optional[Any]=0.1       , lowerCAmelCase__   :		Dict=512       , lowerCAmelCase__   :		Tuple=16       , lowerCAmelCase__   :		Optional[Any]=2       , lowerCAmelCase__   :		Optional[int]=0.02       , lowerCAmelCase__   :		Union[str, Any]=4       , lowerCAmelCase__   :		Any="block_sparse"       , lowerCAmelCase__   :		List[Any]=True       , lowerCAmelCase__   :		Tuple=False       , lowerCAmelCase__   :		Optional[int]=2       , lowerCAmelCase__   :		Optional[int]=3       , )    ->     List[str]:

            '''simple docstring'''




            _UpperCamelCase				    =       parent
            _UpperCamelCase				    =       batch_size
            _UpperCamelCase				    =       seq_length
            _UpperCamelCase				    =       is_training
            _UpperCamelCase				    =       use_attention_mask
            _UpperCamelCase				    =       use_token_type_ids
            _UpperCamelCase				    =       use_labels
            _UpperCamelCase				    =       vocab_size
            _UpperCamelCase				    =       hidden_size
            _UpperCamelCase				    =       num_hidden_layers
            _UpperCamelCase				    =       num_attention_heads
            _UpperCamelCase				    =       intermediate_size
            _UpperCamelCase				    =       hidden_act
            _UpperCamelCase				    =       hidden_dropout_prob
            _UpperCamelCase				    =       attention_probs_dropout_prob
            _UpperCamelCase				    =       max_position_embeddings
            _UpperCamelCase				    =       type_vocab_size
            _UpperCamelCase				    =       type_sequence_label_size
            _UpperCamelCase				    =       initializer_range
            _UpperCamelCase				    =       num_choices

            _UpperCamelCase				    =       rescale_embeddings
            _UpperCamelCase				    =       attention_type
            _UpperCamelCase				    =       use_bias
            _UpperCamelCase				    =       block_size
            _UpperCamelCase				    =       num_random_blocks





       def snake_case__						(							self   :		Optional[Any]     )    ->     Optional[Any]:

            '''simple docstring'''




            _UpperCamelCase				    =       ids_tensor([self.batch_size, self.seq_length]       , self.vocab_size     )

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

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

            _UpperCamelCase				    =       BigBirdConfig(
                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       , attention_type=self.attention_type       , block_size=self.block_size       , num_random_blocks=self.num_random_blocks       , use_bias=self.use_bias       , rescale_embeddings=self.rescale_embeddings       , )

            return config, input_ids, token_type_ids, attention_mask







       def snake_case__						(							self   :		Any     )    ->     Optional[Any]:

            '''simple docstring'''




            _UpperCamelCase				    =       self.prepare_config_and_inputs()
            _UpperCamelCase				    =       config_and_inputs
            _UpperCamelCase				    =       {
                """input_ids""": input_ids,
                """token_type_ids""": token_type_ids,
                """attention_mask""": attention_mask,
            }
            return config, inputs_dict

@require_flax
class 				__lowerCAmelCase     (    lowerCamelCase__						,						unittest.TestCase ):






       """simple docstring"""
       _snake_case	:     List[str]										=    (
           (
               FlaxBigBirdForCausalLM,
               FlaxBigBirdModel,
               FlaxBigBirdForPreTraining,
               FlaxBigBirdForMaskedLM,
               FlaxBigBirdForMultipleChoice,
               FlaxBigBirdForQuestionAnswering,
               FlaxBigBirdForSequenceClassification,
               FlaxBigBirdForTokenClassification,
           )
           if is_flax_available()
           else ()
       )

       _snake_case	:     Dict										=    False
       _snake_case	:     List[str]										=    False





       def snake_case__						(							self   :		Dict     )    ->     str:

            '''simple docstring'''




            _UpperCamelCase				    =       FlaxBigBirdModelTester(self     )





       @slow
       # copied from `test_modeling_flax_common` because it takes much longer than other models
       def snake_case__						(							self   :		int     )    ->     Optional[Any]:

            '''simple docstring'''




            super().test_from_pretrained_save_pretrained()





       @slow
       # copied from `test_modeling_flax_common` because it takes much longer than other models
       def snake_case__						(							self   :		Dict     )    ->     Optional[Any]:

            '''simple docstring'''




            super().test_from_pretrained_with_no_automatic_init()





       @slow
       # copied from `test_modeling_flax_common` because it takes much longer than other models
       def snake_case__						(							self   :		int     )    ->     Tuple:

            '''simple docstring'''




            super().test_no_automatic_init()





       @slow
       # copied from `test_modeling_flax_common` because it takes much longer than other models
       def snake_case__						(							self   :		List[Any]     )    ->     Union[str, Any]:

            '''simple docstring'''




            super().test_hidden_states_output()





       @slow
       def snake_case__						(							self   :		int     )    ->     str:

            '''simple docstring'''




            for model_class_name in self.all_model_classes:
                 _UpperCamelCase				    =       model_class_name.from_pretrained('''google/bigbird-roberta-base'''     )
                 self.assertIsNotNone(lowerCAmelCase__     )





       def snake_case__						(							self   :		int     )    ->     Any:

            '''simple docstring'''




            if self.test_attn_probs:
                 super().test_attention_outputs()





       @slow
       # copied from `test_modeling_flax_common` because it takes much longer than other models
       def snake_case__						(							self   :		int     )    ->     Dict:

            '''simple docstring'''




            _UpperCamelCase				    =       self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                 with self.subTest(model_class.__name__     ):
                      _UpperCamelCase				    =       self._prepare_for_class(lowerCAmelCase__       , lowerCAmelCase__     )
                      _UpperCamelCase				    =       model_class(lowerCAmelCase__     )

                      @jax.jit
                      def model_jitted(lowerCAmelCase__   :		Union[str, Any]       , lowerCAmelCase__   :		Dict=None       , **lowerCAmelCase__   :		Dict     ):
                           return model(input_ids=lowerCAmelCase__       , attention_mask=lowerCAmelCase__       , **lowerCAmelCase__     )

                      with self.subTest('''JIT Enabled'''     ):
                           _UpperCamelCase				    =       model_jitted(**lowerCAmelCase__     ).to_tuple()

                      with self.subTest('''JIT Disabled'''     ):
                           with jax.disable_jit():
                                _UpperCamelCase				    =       model_jitted(**lowerCAmelCase__     ).to_tuple()

                      self.assertEqual(len(lowerCAmelCase__     )       , len(lowerCAmelCase__     )     )
                      for jitted_output, output in zip(lowerCAmelCase__       , lowerCAmelCase__     ):
                           self.assertEqual(jitted_output.shape       , output.shape     )







       def snake_case__						(							self   :		Optional[int]       , lowerCAmelCase__   :		Optional[int]       , lowerCAmelCase__   :		List[Any]       , lowerCAmelCase__   :		Union[str, Any]       , lowerCAmelCase__   :		Dict=1e-5       , lowerCAmelCase__   :		Dict="outputs"       , lowerCAmelCase__   :		Optional[Any]=None     )    ->     Dict:

            '''simple docstring'''




            if name.startswith('''outputs.attentions'''     ):
                 return
            else:
                 super().check_pt_flax_outputs(lowerCAmelCase__       , lowerCAmelCase__       , lowerCAmelCase__       , lowerCAmelCase__       , lowerCAmelCase__       , lowerCAmelCase__     )








'''simple docstring'''




from collections.abc import Iterable
from typing import Any






class 						_UpperCamelCase  :


			'''simple docstring'''

			def __init__(							self     : Optional[int]							,     lowerCAmelCase__     : int | None = None						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		value
					__SCREAMING_SNAKE_CASE       :		Node | None           =		None  # Added in order to delete a node easier
					__SCREAMING_SNAKE_CASE       :		Node | None           =		None
					__SCREAMING_SNAKE_CASE       :		Node | None           =		None



			def __repr__(							self     : Optional[Any]						):





					"""simple docstring"""






					from pprint import pformat

					if self.left is None and self.right is None:
							return str(self.value						)
					return pformat({F"{self.value}": (self.left, self.right)}							,     indent=1						)






class 						_UpperCamelCase  :


			'''simple docstring'''

			def __init__(							self     : Optional[int]							,     lowerCAmelCase__     : Node | None = None						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		str           =		root

			def __str__(							self     : Union[str, Any]						):





					"""simple docstring"""






					return str(self.root						)

			def UpperCamelCase__					(							self     : Dict							,     lowerCAmelCase__     : Node							,     lowerCAmelCase__     : Node | None						):





					"""simple docstring"""






					if new_children is not None:  # reset its kids
							__SCREAMING_SNAKE_CASE       :		List[str]           =		node.parent
					if node.parent is not None:  # reset its parent
							if self.is_right(lowerCAmelCase__						):  # If it is the right children
									__SCREAMING_SNAKE_CASE       :		Any           =		new_children
							else:
									__SCREAMING_SNAKE_CASE       :		int           =		new_children
					else:
							__SCREAMING_SNAKE_CASE       :		int           =		new_children

			def UpperCamelCase__					(							self     : str							,     lowerCAmelCase__     : Node						):





					"""simple docstring"""






					if node.parent and node.parent.right:
							return node == node.parent.right
					return False

			def UpperCamelCase__					(							self     : List[Any]						):





					"""simple docstring"""






					return self.root is None

			def UpperCamelCase__					(							self     : Optional[Any]							,     lowerCAmelCase__     : str						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		Node(lowerCAmelCase__						)  # create a new Node
					if self.empty():  # if Tree is empty
							__SCREAMING_SNAKE_CASE       :		Optional[int]           =		new_node  # set its root
					else:  # Tree is not empty
							__SCREAMING_SNAKE_CASE       :		Optional[int]           =		self.root  # from root
							if parent_node is None:
									return
							while True:  # While we don't get to a leaf
									if value < parent_node.value:  # We go left
											if parent_node.left is None:
													__SCREAMING_SNAKE_CASE       :		List[str]           =		new_node  # We insert the new node in a leaf
													break
											else:
													__SCREAMING_SNAKE_CASE       :		Any           =		parent_node.left
									else:
											if parent_node.right is None:
													__SCREAMING_SNAKE_CASE       :		Tuple           =		new_node
													break
											else:
													__SCREAMING_SNAKE_CASE       :		List[str]           =		parent_node.right
							__SCREAMING_SNAKE_CASE       :		Tuple           =		parent_node

			def UpperCamelCase__					(							self     : str							,     *lowerCAmelCase__     : List[Any]						):





					"""simple docstring"""






					for value in values:
							self.__insert(lowerCAmelCase__						)

			def UpperCamelCase__					(							self     : List[str]							,     lowerCAmelCase__     : Union[str, Any]						):





					"""simple docstring"""






					if self.empty():
							raise IndexError("""Warning: Tree is empty! please use another."""						)
					else:
							__SCREAMING_SNAKE_CASE       :		List[Any]           =		self.root
							# use lazy evaluation here to avoid NoneType Attribute error
							while node is not None and node.value is not value:
									__SCREAMING_SNAKE_CASE       :		Any           =		node.left if value < node.value else node.right
							return node

			def UpperCamelCase__					(							self     : List[Any]							,     lowerCAmelCase__     : Node | None = None						):





					"""simple docstring"""






					if node is None:
							if self.root is None:
									return None
							__SCREAMING_SNAKE_CASE       :		Optional[Any]           =		self.root

					if not self.empty():
							while node.right is not None:
									__SCREAMING_SNAKE_CASE       :		Tuple           =		node.right
					return node

			def UpperCamelCase__					(							self     : Dict							,     lowerCAmelCase__     : Node | None = None						):





					"""simple docstring"""






					if node is None:
							__SCREAMING_SNAKE_CASE       :		Union[str, Any]           =		self.root
					if self.root is None:
							return None
					if not self.empty():
							__SCREAMING_SNAKE_CASE       :		Optional[Any]           =		self.root
							while node.left is not None:
									__SCREAMING_SNAKE_CASE       :		Any           =		node.left
					return node

			def UpperCamelCase__					(							self     : str							,     lowerCAmelCase__     : int						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		List[str]           =		self.search(lowerCAmelCase__						)  # Look for the node with that label
					if node is not None:
							if node.left is None and node.right is None:  # If it has no children
									self.__reassign_nodes(lowerCAmelCase__							,     lowerCAmelCase__						)
							elif node.left is None:  # Has only right children
									self.__reassign_nodes(lowerCAmelCase__							,     node.right						)
							elif node.right is None:  # Has only left children
									self.__reassign_nodes(lowerCAmelCase__							,     node.left						)
							else:
									__SCREAMING_SNAKE_CASE       :		Tuple           =		self.get_max(
									    node.left						)  # Gets the max value of the left branch
									self.remove(tmp_node.value						)  # type: ignore
									__SCREAMING_SNAKE_CASE       :		Optional[Any]           =		(
									    tmp_node.value  # type: ignore
									)  # Assigns the value to the node to delete and keep tree structure

			def UpperCamelCase__					(							self     : Tuple							,     lowerCAmelCase__     : Node | None						):





					"""simple docstring"""






					if node is not None:
							yield node  # Preorder Traversal
							yield from self.preorder_traverse(node.left						)
							yield from self.preorder_traverse(node.right						)

			def UpperCamelCase__					(							self     : Optional[int]							,     lowerCAmelCase__     : Optional[Any]=None						):





					"""simple docstring"""






					if traversal_function is None:
							return self.preorder_traverse(self.root						)
					else:
							return traversal_function(self.root						)

			def UpperCamelCase__					(							self     : str							,     lowerCAmelCase__     : list							,     lowerCAmelCase__     : Node | None						):





					"""simple docstring"""






					if node:
							self.inorder(lowerCAmelCase__							,     node.left						)
							arr.append(node.value						)
							self.inorder(lowerCAmelCase__							,     node.right						)



			def UpperCamelCase__					(							self     : int							,     lowerCAmelCase__     : int							,     lowerCAmelCase__     : Node						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE       :		list[int]           =		[]
					self.inorder(lowerCAmelCase__							,     lowerCAmelCase__						)  # append all values to list using inorder traversal
					return arr[k - 1]


def 				lowerCAmelCase_   (	_lowerCamelCase:   Node | None					):
		__SCREAMING_SNAKE_CASE       :		Optional[Any]           =		[]
		if curr_node is not None:
				__SCREAMING_SNAKE_CASE       :		Optional[int]           =		postorder(curr_node.left					) + postorder(curr_node.right					) + [curr_node]
		return node_list


def 				lowerCAmelCase_   (	):
		__SCREAMING_SNAKE_CASE       :		str           =		(8, 3, 6, 1, 10, 14, 13, 4, 7)
		__SCREAMING_SNAKE_CASE       :		Dict           =		BinarySearchTree()
		for i in testlist:
				t.insert(_lowerCamelCase					)

		# Prints all the elements of the list in order traversal
		print(_lowerCamelCase					)

		if t.search(6					) is not None:
				print("""The value 6 exists"""					)
		else:
				print("""The value 6 doesn't exist"""					)

		if t.search(-1					) is not None:
				print("""The value -1 exists"""					)
		else:
				print("""The value -1 doesn't exist"""					)

		if not t.empty():
				print("""Max Value: """       ,      t.get_max().value					)  # type: ignore
				print("""Min Value: """       ,      t.get_min().value					)  # type: ignore

		for i in testlist:
				t.remove(_lowerCamelCase					)
				print(_lowerCamelCase					)


if __name__ == "__main__":
						import doctest

						doctest.testmod(verbose=True)