






from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


_UpperCamelCase			:	List[Any]								=logging.get_logger(__name__)



def  a__ (__lowercase							:Union[tf.Tensor, np.ndarray]   )		-> List[int]:
				if isinstance(__lowercase					,  np.ndarray   ):
								return list(tensor.shape   )

				_A	:				Tuple   =      tf.shape(__lowercase   )

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

				_A	:				Union[str, Any]   =      tensor.shape.as_list()

				return [dynamic[i] if s is None else s for i, s in enumerate(__lowercase   )]



def  a__ (__lowercase							:tf.Tensor					,  __lowercase							:Optional[int] = None					,  __lowercase							:Optional[str] = None   )		-> tf.Tensor:
				return tf.nn.softmax(logits=logits + 1e-9					,  axis=__lowercase					,  name=__lowercase   )



def  a__ (__lowercase							:Optional[Any]					,  __lowercase							:Optional[Any]					,  __lowercase							:str					,  __lowercase							:Tuple=1e-5					,  __lowercase							:Dict=-1   )		-> str:
				# 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(__lowercase					,  __lowercase   ):
								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
				_A  ,				_A	:				Dict   =      tf.nn.moments(__lowercase					,  axes=[axis]					,  keepdims=__lowercase   )

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

				# Compute layer normalization using the batch_normalization
				# function.
				_A	:				str   =      tf.nn.batch_normalization(
				    __lowercase					,  __lowercase					,  __lowercase					,  offset=__lowercase					,  scale=__lowercase					,  variance_epsilon=__lowercase					,  )
				return outputs



def  a__ (__lowercase							:Any					,  __lowercase							:str=0					,  __lowercase							:List[str]=-1   )		-> Optional[Any]:
				# 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

				_A	:				List[str]   =      tf.shape(__lowercase   )
				_A	:				int   =      tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]   )
				_A	:				Dict   =      tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]					,  axis=0   )
				return tf.reshape(__lowercase					,  __lowercase   )



def  a__ (__lowercase							:tf.Tensor   )		-> tf.Tensor:
				if not isinstance(__lowercase					,  tf.Tensor   ):
								_A	:				Union[str, Any]   =      tf.convert_to_tensor(__lowercase   )  # Catches stray NumPy inputs
				if encoder_attention_mask.shape.rank == 3:
								_A	:				str   =      encoder_attention_mask[:, None, :, :]
				if encoder_attention_mask.shape.rank == 2:
								_A	:				int   =      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))
				_A	:				List[str]   =      (
				    tf.cast(1					,  encoder_attention_mask.dtype   ) - encoder_extended_attention_mask
				) * encoder_extended_attention_mask.dtype.min

				return encoder_extended_attention_mask



def  a__ (__lowercase							:tf.Tensor					,  __lowercase							:int					,  __lowercase							:str = "input_ids"   )		-> None:
				tf.debugging.assert_less(
				    __lowercase					,  tf.cast(__lowercase					,  dtype=tensor.dtype   )					,  message=(
				        f"""The maximum value of {tensor_name} ({tf.math.reduce_max(__lowercase   )}) must be smaller than the embedding """
				        f"""layer's input dimension ({embed_dim}). The likely cause is some problem at tokenization time."""
				    )					,  )



def  a__ (__lowercase							:Dict					,  __lowercase							:Any					,  __lowercase							:Dict   )		-> int:
				_A	:				Union[str, Any]   =      6_4512
				# 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.
				_A	:				Optional[Any]   =      [x for x in data if len(__lowercase   ) > 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}"""   )

				_A	:				Any   =      np.asarray(__lowercase   )

				_A	:				Optional[Any]   =      1
				_A	:				Union[str, Any]   =      np.array_split(__lowercase					,  __lowercase   )

				# 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
								_A	:				List[str]   =      np.array_split(__lowercase					,  __lowercase   )

				if num_chunks > 1:
								for chunk_id, chunk_data in enumerate(__lowercase   ):
												_A	:				Tuple   =      chunk_data
				else:
								_A	:				Optional[int]   =      data



def  a__ (__lowercase							:Optional[Any]					,  __lowercase							:List[str]   )		-> Optional[Any]:
				if name in group.attrs:
								_A	:				List[str]   =      [n.decode('''utf8'''   ) if hasattr(__lowercase					,  '''decode'''   ) else n for n in group.attrs[name]]
				else:
								_A	:				Optional[Any]   =      []
								_A	:				Tuple   =      0
								while "%s%d" % (name, chunk_id) in group.attrs:
												data.extend(
												    [n.decode('''utf8'''   ) if hasattr(__lowercase					,  '''decode'''   ) else n for n in group.attrs['''%s%d''' % (name, chunk_id)]]   )
												chunk_id += 1
				return data






def  a__ (__lowercase							:List[Any]   )		-> Optional[Any]:

				def _expand_single_ad_tensor(__lowercase							:Union[str, Any]   ):
								if isinstance(__lowercase					,  tf.Tensor   ) and t.shape.rank == 1:
												return tf.expand_dims(__lowercase					,  axis=-1   )
								return t

				return tf.nest.map_structure(_expand_single_ad_tensor					,  __lowercase   )











from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_UpperCamelCase			:	Any								={
    'configuration_megatron_bert': ['MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MegatronBertConfig'],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_UpperCamelCase			:	Optional[Any]								=[
	    'MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'MegatronBertForCausalLM',
	    'MegatronBertForMaskedLM',
	    'MegatronBertForMultipleChoice',
	    'MegatronBertForNextSentencePrediction',
	    'MegatronBertForPreTraining',
	    'MegatronBertForQuestionAnswering',
	    'MegatronBertForSequenceClassification',
	    'MegatronBertForTokenClassification',
	    'MegatronBertModel',
	    'MegatronBertPreTrainedModel',
	]

if TYPE_CHECKING:
	from .configuration_megatron_bert import MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, MegatronBertConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_megatron_bert import (
		    MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    MegatronBertForCausalLM,
		    MegatronBertForMaskedLM,
		    MegatronBertForMultipleChoice,
		    MegatronBertForNextSentencePrediction,
		    MegatronBertForPreTraining,
		    MegatronBertForQuestionAnswering,
		    MegatronBertForSequenceClassification,
		    MegatronBertForTokenClassification,
		    MegatronBertModel,
		    MegatronBertPreTrainedModel,
		)

else:
	import sys

	_UpperCamelCase			:	Optional[int]								=_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



