




"""simple docstring"""




from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


a_        =       logging.get_logger(__name__)







def      __UpperCAmelCase     ( __UpperCamelCase	):
			if isinstance(__snake_case	,				np.ndarray	):
						return list(tensor.shape	)

			__lowercase			:  Optional[Any]								=       tf.shape(__snake_case	)

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

			__lowercase			:  Dict								=       tensor.shape.as_list()

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







def      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase = None	,				__UpperCamelCase = None	):
			return tf.nn.softmax(logits=logits + 1e-9	,				axis=__snake_case	,				name=__snake_case	)







def      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase=1e-5	,				__UpperCamelCase=-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(__snake_case	,				__snake_case	):
						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			:  Optional[int]								=       tf.nn.moments(__snake_case	,				axes=[axis]	,				keepdims=__snake_case	)

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

			# Compute layer normalization using the batch_normalization
			# function.
			__lowercase			:  List[Any]								=       tf.nn.batch_normalization(
			    __snake_case	,				__snake_case	,				__snake_case	,				offset=__snake_case	,				scale=__snake_case	,				variance_epsilon=__snake_case	,				)
			return outputs







def      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase=0	,				__UpperCamelCase=-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			:  str								=       tf.shape(__snake_case	)
			__lowercase			:  Any								=       tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]	)
			__lowercase			:  Optional[int]								=       tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]	,				axis=0	)
			return tf.reshape(__snake_case	,				__snake_case	)







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

			return encoder_extended_attention_mask







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







def      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	):
			__lowercase			:  List[Any]								=       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			:  Union[str, Any]								=       [x for x in data if len(__snake_case	) > 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			:  Dict								=       np.asarray(__snake_case	)

			__lowercase			:  Optional[int]								=       1
			__lowercase			:  int								=       np.array_split(__snake_case	,				__snake_case	)

			# 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[int]								=       np.array_split(__snake_case	,				__snake_case	)

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







def      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase	):
			if name in group.attrs:
						__lowercase			:  Any								=       [n.decode('''utf8'''	) if hasattr(__snake_case	,				'''decode'''	) else n for n in group.attrs[name]]
			else:
						__lowercase			:  Union[str, Any]								=       []
						__lowercase			:  Dict								=       0
						while "%s%d" % (name, chunk_id) in group.attrs:
									data.extend(
									    [n.decode('''utf8'''	) if hasattr(__snake_case	,				'''decode'''	) else n for n in group.attrs['''%s%d''' % (name, chunk_id)]]	)
									chunk_id += 1
			return data







def      __UpperCAmelCase     ( __UpperCamelCase	):

			def _expand_single_ad_tensor(__UpperCamelCase	):
						if isinstance(__snake_case	,				tf.Tensor	) and t.shape.rank == 1:
									return tf.expand_dims(__snake_case	,				axis=-1	)
						return t

			return tf.nest.map_structure(_expand_single_ad_tensor	,				__snake_case	)




import logging
import os

from .state import PartialState
class     lowerCAmelCase_     (       logging.LoggerAdapter ):







							@staticmethod
							def  UpperCamelCase_						(   _A      :     Any		):
												_UpperCamelCase			     =     PartialState()
												return not main_process_only or (main_process_only and state.is_main_process)







							def  UpperCamelCase_						(   self      :     Union[str, Any] ,			_A      :     Optional[Any] ,			_A      :     str ,			*_A      :     int ,			**_A      :     List[Any]		):
												if PartialState._shared_state == {}:
																	raise RuntimeError(
																	    '''You must initialize the accelerate state by calling either `PartialState()` or `Accelerator()` before using the logging utility.'''		)
												_UpperCamelCase			     =     kwargs.pop('''main_process_only''' ,			_A		)
												_UpperCamelCase			     =     kwargs.pop('''in_order''' ,			_A		)

												if self.isEnabledFor(_A		):
																	if self._should_log(_A		):
																						_UpperCamelCase     ,    _UpperCamelCase			     =     self.process(_A ,			_A		)
																						self.logger.log(_A ,			_A ,			*_A ,			**_A		)

																	elif in_order:
																						_UpperCamelCase			     =     PartialState()
																						for i in range(state.num_processes		):
																											if i == state.process_index:
																																_UpperCamelCase     ,    _UpperCamelCase			     =     self.process(_A ,			_A		)
																																self.logger.log(_A ,			_A ,			*_A ,			**_A		)
																											state.wait_for_everyone()




def 					_snake_case	(  __snake_case  ,     __snake_case = None						):
					if log_level is None:
										_UpperCamelCase			     =     os.environ.get('''ACCELERATE_LOG_LEVEL'''  ,     __snake_case						)
					_UpperCamelCase			     =     logging.getLogger(__snake_case						)
					if log_level is not None:
										logger.setLevel(log_level.upper()						)
										logger.root.setLevel(log_level.upper()						)
					return MultiProcessAdapter(__snake_case  ,     {}						)
