





'''simple docstring'''


from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


lowerCAmelCase					:     str												=logging.get_logger(__name__)
def 							UpperCAmelCase_  (						__lowerCamelCase					:					Union[tf.Tensor, np.ndarray] ):
   if isinstance(__lowerCamelCase	,np.ndarray ):
      return list(tensor.shape )

   lowercase_			:Optional[Any]    			= tf.shape(__lowerCamelCase )

   if tensor.shape == tf.TensorShape(__lowerCamelCase ):
      return dynamic

   lowercase_			:Optional[Any]    			= tensor.shape.as_list()

   return [dynamic[i] if s is None else s for i, s in enumerate(__lowerCamelCase )]
def 							UpperCAmelCase_  (						__lowerCamelCase					:					tf.Tensor	,__lowerCamelCase					:					Optional[int] = None	,__lowerCamelCase					:					Optional[str] = None ):
   return tf.nn.softmax(logits=logits + 1e-9	,axis=__lowerCamelCase	,name=__lowerCamelCase )
def 							UpperCAmelCase_  (						__lowerCamelCase					:					Optional[int]	,__lowerCamelCase					:					str	,__lowerCamelCase					:					Any	,__lowerCamelCase					:					List[Any]=1e-5	,__lowerCamelCase					:					Dict=-1 ):
   # This is a very simplified functional layernorm, designed to duplicate
   # the functionality of PyTorch nn.functional.layer_norm when this is needed to port
   # models in Transformers.

   if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(__lowerCamelCase	,__lowerCamelCase ):
      raise NotImplementedError("Only 1D weight and bias tensors are supported for now, with only a single axis." )

   # Get mean and variance on the axis to be normalized
   lowercase_			,							lowercase_			:List[Any]    			= tf.nn.moments(__lowerCamelCase	,axes=[axis]	,keepdims=__lowerCamelCase )

   if axis != -1:
      # Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
      # on every dimension except axis
      lowercase_			:List[Any]    			= [1] * inputs.shape.rank
      lowercase_			:Optional[Any]    			= shape_list(__lowerCamelCase )[axis]
      lowercase_			:Tuple    			= tf.reshape(__lowerCamelCase	,__lowerCamelCase )
      lowercase_			:Tuple    			= tf.reshape(__lowerCamelCase	,__lowerCamelCase )

   # Compute layer normalization using the batch_normalization
   # function.
   lowercase_			:Optional[Any]    			= tf.nn.batch_normalization(
       __lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,offset=__lowerCamelCase	,scale=__lowerCamelCase	,variance_epsilon=__lowerCamelCase	,)
   return outputs
def 							UpperCAmelCase_  (						__lowerCamelCase					:					List[Any]	,__lowerCamelCase					:					Any=0	,__lowerCamelCase					:					Tuple=-1 ):
   # Replicates the behavior of torch.flatten in TF

   # If end_dim or start_dim is negative, count them from the end
   if end_dim < 0:
      end_dim += input.shape.rank
   if start_dim < 0:
      start_dim += input.shape.rank

   if start_dim == end_dim:
      return input

   lowercase_			:Tuple    			= tf.shape(__lowerCamelCase )
   lowercase_			:List[Any]    			= tf.math.reduce_prod(in_shape[start_dim : end_dim + 1] )
   lowercase_			:str    			= tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]	,axis=0 )
   return tf.reshape(__lowerCamelCase	,__lowerCamelCase )
def 							UpperCAmelCase_  (						__lowerCamelCase					:					tf.Tensor ):
   if not isinstance(__lowerCamelCase	,tf.Tensor ):
      lowercase_			:Tuple    			= tf.convert_to_tensor(__lowerCamelCase )  # Catches stray NumPy inputs
   if encoder_attention_mask.shape.rank == 3:
      lowercase_			:List[str]    			= encoder_attention_mask[:, None, :, :]
   if encoder_attention_mask.shape.rank == 2:
      lowercase_			:Tuple    			= encoder_attention_mask[:, None, None, :]
   # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
   # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
   # /transformer/transformer_layers.py#L270
   # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
   # encoder_extended_attention_mask.transpose(-1, -2))
   lowercase_			:Tuple    			= (
       tf.cast(1	,encoder_attention_mask.dtype ) - encoder_extended_attention_mask
   ) * encoder_extended_attention_mask.dtype.min

   return encoder_extended_attention_mask
def 							UpperCAmelCase_  (						__lowerCamelCase					:					tf.Tensor	,__lowerCamelCase					:					int	,__lowerCamelCase					:					str = "input_ids" ):
   tf.debugging.assert_less(
       __lowerCamelCase	,tf.cast(__lowerCamelCase	,dtype=tensor.dtype )	,message=(
           F'The maximum value of {tensor_name} ({tf.math.reduce_max(__lowerCamelCase )}) must be smaller than the embedding '
           F'layer\'s input dimension ({embed_dim}). The likely cause is some problem at tokenization time.'
       )	,)
def 							UpperCAmelCase_  (						__lowerCamelCase					:					Dict	,__lowerCamelCase					:					List[str]	,__lowerCamelCase					:					Optional[int] ):
   lowercase_			:int    			= 6_45_12
   # Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
   # because in that case even chunking the array would not make the saving
   # possible.
   lowercase_			:Any    			= [x for x in data if len(__lowerCamelCase ) > HDF5_OBJECT_HEADER_LIMIT]

   # Expecting this to never be true.
   if bad_attributes:
      raise RuntimeError(
          "The following attributes cannot be saved to HDF5 file because "
          F'they are larger than {HDF5_OBJECT_HEADER_LIMIT} '
          F'bytes: {bad_attributes}' )

   lowercase_			:Optional[int]    			= np.asarray(__lowerCamelCase )

   lowercase_			:Optional[int]    			= 1
   lowercase_			:Optional[Any]    			= np.array_split(__lowerCamelCase	,__lowerCamelCase )

   # This will never loop forever thanks to the test above.
   while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data ):
      num_chunks += 1
      lowercase_			:Optional[Any]    			= np.array_split(__lowerCamelCase	,__lowerCamelCase )

   if num_chunks > 1:
      for chunk_id, chunk_data in enumerate(__lowerCamelCase ):
         lowercase_			:Optional[int]    			= chunk_data
   else:
      lowercase_			:Dict    			= data
def 							UpperCAmelCase_  (						__lowerCamelCase					:					Optional[Any]	,__lowerCamelCase					:					Union[str, Any] ):
   if name in group.attrs:
      lowercase_			:Tuple    			= [n.decode("utf8" ) if hasattr(__lowerCamelCase	,"decode" ) else n for n in group.attrs[name]]
   else:
      lowercase_			:Union[str, Any]    			= []
      lowercase_			:Optional[Any]    			= 0
      while "%s%d" % (name, chunk_id) in group.attrs:
         data.extend(
             [n.decode("utf8" ) if hasattr(__lowerCamelCase	,"decode" ) else n for n in group.attrs["%s%d" % (name, chunk_id)]] )
         chunk_id += 1
   return data


def 							UpperCAmelCase_  (						__lowerCamelCase					:					Union[str, Any] ):

   def _expand_single_ad_tensor(__lowerCamelCase					:					Optional[int] ):
      if isinstance(__lowerCamelCase	,tf.Tensor ) and t.shape.rank == 1:
         return tf.expand_dims(__lowerCamelCase	,axis=-1 )
      return t

   return tf.nest.map_structure(_expand_single_ad_tensor	,__lowerCamelCase )









'''simple docstring'''


import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
 from transformers.models.bert.modeling_flax_bert import (
     FlaxBertForMaskedLM,
     FlaxBertForMultipleChoice,
     FlaxBertForNextSentencePrediction,
     FlaxBertForPreTraining,
     FlaxBertForQuestionAnswering,
     FlaxBertForSequenceClassification,
     FlaxBertForTokenClassification,
     FlaxBertModel,
 )
class       a_   (				unittest.TestCase							):





     def __init__( self		:					List[str]	,      lowercase		:					str	,      lowercase		:					Union[str, Any]=13	,      lowercase		:					int=7	,      lowercase		:					List[str]=True	,      lowercase		:					int=True	,      lowercase		:					str=True	,      lowercase		:					Any=True	,      lowercase		:					List[str]=99	,      lowercase		:					Union[str, Any]=32	,      lowercase		:					Optional[Any]=5	,      lowercase		:					Dict=4	,      lowercase		:					Dict=37	,      lowercase		:					Dict="gelu"	,      lowercase		:					Optional[int]=0.1	,      lowercase		:					str=0.1	,      lowercase		:					List[Any]=512	,      lowercase		:					str=16	,      lowercase		:					Dict=2	,      lowercase		:					Any=0.02	,      lowercase		:					Any=4	,      ):





        """simple docstring"""




        lowercase_			:List[str]    			= parent
        lowercase_			:Any    			= batch_size
        lowercase_			:Dict    			= seq_length
        lowercase_			:Union[str, Any]    			= is_training
        lowercase_			:Optional[int]    			= use_attention_mask
        lowercase_			:Any    			= use_token_type_ids
        lowercase_			:Union[str, Any]    			= use_labels
        lowercase_			:Dict    			= vocab_size
        lowercase_			:Tuple    			= hidden_size
        lowercase_			:Tuple    			= num_hidden_layers
        lowercase_			:Optional[int]    			= num_attention_heads
        lowercase_			:Optional[Any]    			= intermediate_size
        lowercase_			:str    			= hidden_act
        lowercase_			:Tuple    			= hidden_dropout_prob
        lowercase_			:Optional[Any]    			= attention_probs_dropout_prob
        lowercase_			:Tuple    			= max_position_embeddings
        lowercase_			:Any    			= type_vocab_size
        lowercase_			:int    			= type_sequence_label_size
        lowercase_			:Tuple    			= initializer_range
        lowercase_			:Optional[Any]    			= num_choices





     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:int    			= ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size						)

        lowercase_			:Union[str, Any]    			= None
        if self.use_attention_mask:
           lowercase_			:Dict    			= random_attention_mask([self.batch_size, self.seq_length]						)

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

        lowercase_			:Optional[Any]    			= BertConfig(
            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=lowercase	,      initializer_range=self.initializer_range	,      )

        return config, input_ids, token_type_ids, attention_mask





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





        """simple docstring"""




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







     def 							lowercase__      ( self		:					List[Any]						):





        """simple docstring"""




        lowercase_			:Any    			= self.prepare_config_and_inputs()
        lowercase_			,							lowercase_			,							lowercase_			,							lowercase_			:Union[str, Any]    			= config_and_inputs

        lowercase_			:Dict    			= True
        lowercase_			:Optional[int]    			= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
        lowercase_			:str    			= ids_tensor([self.batch_size, self.seq_length]	,      vocab_size=2						)

        return (
            config,
            input_ids,
            attention_mask,
            encoder_hidden_states,
            encoder_attention_mask,
        )







@require_flax
class       a_   (				_lowerCAmelCase	,						unittest.TestCase							):
     __A      =			True

     __A      =			(
         (
             FlaxBertModel,
             FlaxBertForPreTraining,
             FlaxBertForMaskedLM,
             FlaxBertForMultipleChoice,
             FlaxBertForQuestionAnswering,
             FlaxBertForNextSentencePrediction,
             FlaxBertForSequenceClassification,
             FlaxBertForTokenClassification,
             FlaxBertForQuestionAnswering,
         )
         if is_flax_available()
         else ()
     )





     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= FlaxBertModelTester(self						)







     @slow
     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        lowercase_			:List[str]    			= FlaxBertModel.from_pretrained("bert-base-cased"						)
        lowercase_			:str    			= model(np.ones((1, 1)						)						)
        self.assertIsNotNone(lowercase						)


