





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






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase=2					,	__lowerCAmelCase=56					,	__lowerCAmelCase=True					,	__lowerCAmelCase=True					,	__lowerCAmelCase=True					,	__lowerCAmelCase=True					,	__lowerCAmelCase=99					,	__lowerCAmelCase=32					,	__lowerCAmelCase=2					,	__lowerCAmelCase=2					,	__lowerCAmelCase=7					,	__lowerCAmelCase="gelu_new"					,	__lowerCAmelCase=0.1					,	__lowerCAmelCase=0.1					,	__lowerCAmelCase=512					,	__lowerCAmelCase=16					,	__lowerCAmelCase=2					,	__lowerCAmelCase=0.02					,	__lowerCAmelCase=4					,	__lowerCAmelCase="block_sparse"					,	__lowerCAmelCase=True					,	__lowerCAmelCase=False					,	__lowerCAmelCase=2					,	__lowerCAmelCase=3					,	):
														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 			_lowerCamelCase						(  self					):
														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 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     self.prepare_config_and_inputs()
														UpperCamelCase__       ,       UpperCamelCase__       ,       UpperCamelCase__       ,       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     __SCREAMING_SNAKE_CASE   (       _a				,						unittest.TestCase    ):
							snake_case       : List[Any]         =    (
							    (
							        FlaxBigBirdForCausalLM,
							        FlaxBigBirdModel,
							        FlaxBigBirdForPreTraining,
							        FlaxBigBirdForMaskedLM,
							        FlaxBigBirdForMultipleChoice,
							        FlaxBigBirdForQuestionAnswering,
							        FlaxBigBirdForSequenceClassification,
							        FlaxBigBirdForTokenClassification,
							    )
							    if is_flax_available()
							    else ()
							)

							snake_case       : List[Any]         =    False
							snake_case       : Any         =    False






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     FlaxBigBirdModelTester(self					)






							@slow
							# copied from `test_modeling_flax_common` because it takes much longer than other models
							def 			_lowerCamelCase						(  self					):
														super().test_from_pretrained_save_pretrained()






							@slow
							# copied from `test_modeling_flax_common` because it takes much longer than other models
							def 			_lowerCamelCase						(  self					):
														super().test_from_pretrained_with_no_automatic_init()






							@slow
							# copied from `test_modeling_flax_common` because it takes much longer than other models
							def 			_lowerCamelCase						(  self					):
														super().test_no_automatic_init()






							@slow
							# copied from `test_modeling_flax_common` because it takes much longer than other models
							def 			_lowerCamelCase						(  self					):
														super().test_hidden_states_output()






							@slow
							def 			_lowerCamelCase						(  self					):
														for model_class_name in self.all_model_classes:
																					UpperCamelCase__			       =     model_class_name.from_pretrained("""google/bigbird-roberta-base"""					)
																					self.assertIsNotNone(__lowerCAmelCase					)






							def 			_lowerCamelCase						(  self					):
														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 			_lowerCamelCase						(  self					):
														UpperCamelCase__       ,       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					,	__lowerCAmelCase=None					,	**__lowerCAmelCase					):
																																			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 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase=1E-5					,	__lowerCAmelCase="outputs"					,	__lowerCAmelCase=None					):
														# `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
														# an effort was done to return `attention_probs` (yet to be verified).
														if name.startswith("""outputs.attentions"""					):
																					return
														else:
																					super().check_pt_flax_outputs(__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					,	__lowerCAmelCase					)











from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map





def       _UpperCamelCase       (a__				:Dict[str, torch.Tensor]    ):






							"""simple docstring"""





							UpperCamelCase__			       =     []
							UpperCamelCase__			       =     []
							UpperCamelCase__			       =     []

							for rt in rc.restypes:
														UpperCamelCase__			       =     rc.restype_name_to_atomaa_names[rc.restype_atoa[rt]]
														restype_atomaa_to_atomaa_list.append([(rc.atom_order[name] if name else 0) for name in atom_names]    )
														UpperCamelCase__			       =     {name: i for i, name in enumerate(a__    )}
														restype_atomaa_to_atomaa_list.append(
														    [(atom_name_to_idxaa[name] if name in atom_name_to_idxaa else 0) for name in rc.atom_types]    )

														restype_atomaa_mask_list.append([(1.0 if name else 0.0) for name in atom_names]    )

							# Add dummy mapping for restype 'UNK'
							restype_atomaa_to_atomaa_list.append([0] * 14    )
							restype_atomaa_to_atomaa_list.append([0] * 37    )
							restype_atomaa_mask_list.append([0.0] * 14    )

							UpperCamelCase__			       =     torch.tensor(
							    a__				,     dtype=torch.intaa				,     device=protein["""aatype"""].device				,     )
							UpperCamelCase__			       =     torch.tensor(
							    a__				,     dtype=torch.intaa				,     device=protein["""aatype"""].device				,     )
							UpperCamelCase__			       =     torch.tensor(
							    a__				,     dtype=torch.floataa				,     device=protein["""aatype"""].device				,     )
							UpperCamelCase__			       =     protein["""aatype"""].to(torch.long    )

							# create the mapping for (residx, atom14) --> atom37, i.e. an array
							# with shape (num_res, 14) containing the atom37 indices for this protein
							UpperCamelCase__			       =     restype_atomaa_to_atomaa[protein_aatype]
							UpperCamelCase__			       =     restype_atomaa_mask[protein_aatype]

							UpperCamelCase__			       =     residx_atomaa_mask
							UpperCamelCase__			       =     residx_atomaa_to_atomaa.long()

							# create the gather indices for mapping back
							UpperCamelCase__			       =     restype_atomaa_to_atomaa[protein_aatype]
							UpperCamelCase__			       =     residx_atomaa_to_atomaa.long()

							# create the corresponding mask
							UpperCamelCase__			       =     torch.zeros([21, 37]				,     dtype=torch.floataa				,     device=protein["""aatype"""].device    )
							for restype, restype_letter in enumerate(rc.restypes    ):
														UpperCamelCase__			       =     rc.restype_atoa[restype_letter]
														UpperCamelCase__			       =     rc.residue_atoms[restype_name]
														for atom_name in atom_names:
																					UpperCamelCase__			       =     rc.atom_order[atom_name]
																					UpperCamelCase__			       =     1

							UpperCamelCase__			       =     restype_atomaa_mask[protein_aatype]
							UpperCamelCase__			       =     residx_atomaa_mask

							return protein





def       _UpperCamelCase       (a__				:Dict[str, torch.Tensor]    ):






							"""simple docstring"""





							UpperCamelCase__			       =     tree_map(lambda a__    : torch.tensor(a__				,     device=batch["""aatype"""].device    )				,     a__				,     np.ndarray    )
							UpperCamelCase__			       =     tensor_tree_map(lambda a__    : np.array(a__    )				,     make_atomaa_masks(a__    )    )
							return out




