





"""simple docstring"""





import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
							import torch
							from torch import nn

							from transformers import (
							    Adafactor,
							    AdamW,
							    get_constant_schedule,
							    get_constant_schedule_with_warmup,
							    get_cosine_schedule_with_warmup,
							    get_cosine_with_hard_restarts_schedule_with_warmup,
							    get_inverse_sqrt_schedule,
							    get_linear_schedule_with_warmup,
							    get_polynomial_decay_schedule_with_warmup,
							)



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Optional[int]    ,						SCREAMING_SNAKE_CASE       :							List[str]=10   ):


							'''simple docstring'''







							__lowerCamelCase      :      Any         =[]
							for _ in range(__snake_case   ):
														lrs.append(scheduler.get_lr()[0]   )
														scheduler.step()
							return lrs



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							str    ,						SCREAMING_SNAKE_CASE       :							Union[str, Any]=10   ):


							'''simple docstring'''







							__lowerCamelCase      :      List[Any]         =[]
							for step in range(__snake_case   ):
														lrs.append(scheduler.get_lr()[0]   )
														scheduler.step()
														if step == num_steps // 2:
																					with tempfile.TemporaryDirectory() as tmpdirname:
																												__lowerCamelCase      :      Union[str, Any]         =os.path.join(__snake_case    ,						'''schedule.bin'''   )
																												torch.save(scheduler.state_dict()    ,						__snake_case   )

																												__lowerCamelCase      :      List[Any]         =torch.load(__snake_case   )
																												scheduler.load_state_dict(__snake_case   )
							return lrs
@require_torch
class 		SCREAMING_SNAKE_CASE_							(    unittest.TestCase   ):
				"""simple docstring"""




				def 					__lowercase  (     self  :Dict						,    __lowercase  :Optional[Any]						,    __lowercase  :Optional[Any]						,    __lowercase  :Tuple			):
											self.assertEqual(len(__lowercase			)						,    len(__lowercase			)			)
											for a, b in zip(__lowercase						,    __lowercase			):
																		self.assertAlmostEqual(__lowercase						,    __lowercase						,    delta=__lowercase			)




				def 					__lowercase  (     self  :Any			):
											__lowerCamelCase      :      Union[str, Any]         =torch.tensor([0.1, -0.2, -0.1]						,    requires_grad=__lowercase			)
											__lowerCamelCase      :      List[Any]         =torch.tensor([0.4, 0.2, -0.5]			)
											__lowerCamelCase      :      int         =nn.MSELoss()
											# No warmup, constant schedule, no gradient clipping
											__lowerCamelCase      :      List[str]         =AdamW(params=[w]						,    lr=2e-1						,    weight_decay=0.0			)
											for _ in range(100			):
																		__lowerCamelCase      :      str         =criterion(__lowercase						,    __lowercase			)
																		loss.backward()
																		optimizer.step()
																		w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
																		w.grad.zero_()
											self.assertListAlmostEqual(w.tolist()						,    [0.4, 0.2, -0.5]						,    tol=1e-2			)




				def 					__lowercase  (     self  :Union[str, Any]			):
											__lowerCamelCase      :      int         =torch.tensor([0.1, -0.2, -0.1]						,    requires_grad=__lowercase			)
											__lowerCamelCase      :      Optional[Any]         =torch.tensor([0.4, 0.2, -0.5]			)
											__lowerCamelCase      :      Optional[int]         =nn.MSELoss()
											# No warmup, constant schedule, no gradient clipping
											__lowerCamelCase      :      Any         =Adafactor(
											    params=[w]						,    lr=1e-2						,    eps=(1e-3_0, 1e-3)						,    clip_threshold=1.0						,    decay_rate=-0.8						,    betaa=__lowercase						,    weight_decay=0.0						,    relative_step=__lowercase						,    scale_parameter=__lowercase						,    warmup_init=__lowercase						,    )
											for _ in range(1000			):
																		__lowerCamelCase      :      Any         =criterion(__lowercase						,    __lowercase			)
																		loss.backward()
																		optimizer.step()
																		w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
																		w.grad.zero_()
											self.assertListAlmostEqual(w.tolist()						,    [0.4, 0.2, -0.5]						,    tol=1e-2			)

@require_torch
class 		SCREAMING_SNAKE_CASE_							(    unittest.TestCase   ):
				"""simple docstring"""
				__snake_case	 =		nn.Linear(50       ,  50   ) if is_torch_available() else None
				__snake_case	 =		AdamW(m.parameters()       ,  lr=10.0   ) if is_torch_available() else None
				__snake_case	 =		10




				def 					__lowercase  (     self  :Dict						,    __lowercase  :Dict						,    __lowercase  :str						,    __lowercase  :List[Any]						,    __lowercase  :Any=None			):
											self.assertEqual(len(__lowercase			)						,    len(__lowercase			)			)
											for a, b in zip(__lowercase						,    __lowercase			):
																		self.assertAlmostEqual(__lowercase						,    __lowercase						,    delta=__lowercase						,    msg=__lowercase			)




				def 					__lowercase  (     self  :Union[str, Any]			):
											__lowerCamelCase      :      Dict         ={"num_warmup_steps": 2, "num_training_steps": 10}
											# schedulers doct format
											# function: (sched_args_dict, expected_learning_rates)
											__lowerCamelCase      :      Dict         ={
											    get_constant_schedule: ({}, [10.0] * self.num_steps),
											    get_constant_schedule_with_warmup: (
											        {"num_warmup_steps": 4},
											        [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
											    ),
											    get_linear_schedule_with_warmup: (
											        {**common_kwargs},
											        [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
											    ),
											    get_cosine_schedule_with_warmup: (
											        {**common_kwargs},
											        [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
											    ),
											    get_cosine_with_hard_restarts_schedule_with_warmup: (
											        {**common_kwargs, "num_cycles": 2},
											        [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
											    ),
											    get_polynomial_decay_schedule_with_warmup: (
											        {**common_kwargs, "power": 2.0, "lr_end": 1e-7},
											        [0.0, 5.0, 10.0, 7.656, 5.625, 3.906, 2.5, 1.406, 0.625, 0.156],
											    ),
											    get_inverse_sqrt_schedule: (
											        {"num_warmup_steps": 2},
											        [0.0, 5.0, 10.0, 8.165, 7.071, 6.325, 5.774, 5.345, 5.0, 4.714],
											    ),
											}

											for scheduler_func, data in scheds.items():
																		__lowerCamelCase      :      List[Any]         =data

																		__lowerCamelCase      :      List[str]         =scheduler_func(self.optimizer						,    **__lowercase			)
																		self.assertEqual(len([scheduler.get_lr()[0]]			)						,    1			)
																		__lowerCamelCase      :      int         =unwrap_schedule(__lowercase						,    self.num_steps			)
																		self.assertListAlmostEqual(
																		    __lowercase						,    __lowercase						,    tol=1e-2						,    msg=f'failed for {scheduler_func} in normal scheduler'						,    )

																		__lowerCamelCase      :      str         =scheduler_func(self.optimizer						,    **__lowercase			)
																		if scheduler_func.__name__ != "get_constant_schedule":
																									LambdaScheduleWrapper.wrap_scheduler(__lowercase			)  # wrap to test picklability of the schedule
																		__lowerCamelCase      :      Dict         =unwrap_and_save_reload_schedule(__lowercase						,    self.num_steps			)
																		self.assertListEqual(__lowercase						,    __lowercase						,    msg=f'failed for {scheduler_func} in save and reload'			)

class 		SCREAMING_SNAKE_CASE_							:
				"""simple docstring"""




				def __init__(     self  :Any						,    __lowercase  :Any			):
											__lowerCamelCase      :      Any         =fn




				def __call__(     self  :Optional[int]						,    *__lowercase  :List[str]						,    **__lowercase  :Optional[Any]			):
											return self.fn(*__lowercase						,    **__lowercase			)




				@classmethod
				def 					__lowercase  (     self  :Optional[int]						,    __lowercase  :Dict			):
											__lowerCamelCase      :      Dict         =list(map(self						,    scheduler.lr_lambdas			)			)











"""simple docstring"""





from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


_UpperCamelCase							=      logging.get_logger(__name__)



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Union[tf.Tensor, np.ndarray]   ):


							'''simple docstring'''







							if isinstance(SCREAMING_SNAKE_CASE    ,						np.ndarray   ):
														return list(tensor.shape   )

							__lowerCamelCase      :      Tuple         =tf.shape(SCREAMING_SNAKE_CASE   )

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

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

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



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							tf.Tensor    ,						SCREAMING_SNAKE_CASE       :							Optional[int] = None    ,						SCREAMING_SNAKE_CASE       :							Optional[str] = None   ):


							'''simple docstring'''







							return tf.nn.softmax(logits=logits + 1E-9    ,						axis=SCREAMING_SNAKE_CASE    ,						name=SCREAMING_SNAKE_CASE   )



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Optional[int]    ,						SCREAMING_SNAKE_CASE       :							int    ,						SCREAMING_SNAKE_CASE       :							Dict    ,						SCREAMING_SNAKE_CASE       :							Tuple=1E-5    ,						SCREAMING_SNAKE_CASE       :							Union[str, Any]=-1   ):


							'''simple docstring'''







							if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(SCREAMING_SNAKE_CASE    ,						SCREAMING_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
							__lowerCamelCase    ,						__lowerCamelCase      :      Dict         =tf.nn.moments(SCREAMING_SNAKE_CASE    ,						axes=[axis]    ,						keepdims=SCREAMING_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
														__lowerCamelCase      :      Dict         =[1] * inputs.shape.rank
														__lowerCamelCase      :      Dict         =shape_list(SCREAMING_SNAKE_CASE   )[axis]
														__lowerCamelCase      :      List[str]         =tf.reshape(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )
														__lowerCamelCase      :      List[str]         =tf.reshape(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )

							# Compute layer normalization using the batch_normalization
							# function.
							__lowerCamelCase      :      int         =tf.nn.batch_normalization(
							    SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE    ,						offset=SCREAMING_SNAKE_CASE    ,						scale=SCREAMING_SNAKE_CASE    ,						variance_epsilon=SCREAMING_SNAKE_CASE    ,						)
							return outputs



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							List[Any]    ,						SCREAMING_SNAKE_CASE       :							Dict=0    ,						SCREAMING_SNAKE_CASE       :							int=-1   ):


							'''simple docstring'''







							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

							__lowerCamelCase      :      List[str]         =tf.shape(SCREAMING_SNAKE_CASE   )
							__lowerCamelCase      :      Optional[Any]         =tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]   )
							__lowerCamelCase      :      int         =tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]    ,						axis=0   )
							return tf.reshape(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							tf.Tensor   ):


							'''simple docstring'''







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

							return encoder_extended_attention_mask



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							tf.Tensor    ,						SCREAMING_SNAKE_CASE       :							int    ,						SCREAMING_SNAKE_CASE       :							str = "input_ids"   ):


							'''simple docstring'''







							tf.debugging.assert_less(
							    SCREAMING_SNAKE_CASE    ,						tf.cast(SCREAMING_SNAKE_CASE    ,						dtype=tensor.dtype   )    ,						message=(
							        F'The maximum value of {tensor_name} ({tf.math.reduce_max(SCREAMING_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 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							str    ,						SCREAMING_SNAKE_CASE       :							str    ,						SCREAMING_SNAKE_CASE       :							str   ):


							'''simple docstring'''







							__lowerCamelCase      :      Optional[int]         =64512
							# 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.
							__lowerCamelCase      :      Any         =[x for x in data if len(SCREAMING_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}'   )

							__lowerCamelCase      :      Tuple         =np.asarray(SCREAMING_SNAKE_CASE   )

							__lowerCamelCase      :      Any         =1
							__lowerCamelCase      :      Optional[Any]         =np.array_split(SCREAMING_SNAKE_CASE    ,						SCREAMING_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
														__lowerCamelCase      :      Dict         =np.array_split(SCREAMING_SNAKE_CASE    ,						SCREAMING_SNAKE_CASE   )

							if num_chunks > 1:
														for chunk_id, chunk_data in enumerate(SCREAMING_SNAKE_CASE   ):
																					__lowerCamelCase      :      List[Any]         =chunk_data
							else:
														__lowerCamelCase      :      Optional[int]         =data



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Optional[int]    ,						SCREAMING_SNAKE_CASE       :							Optional[int]   ):


							'''simple docstring'''







							if name in group.attrs:
														__lowerCamelCase      :      Optional[int]         =[n.decode('''utf8'''   ) if hasattr(SCREAMING_SNAKE_CASE    ,						'''decode'''   ) else n for n in group.attrs[name]]
							else:
														__lowerCamelCase      :      Tuple         =[]
														__lowerCamelCase      :      List[str]         =0
														while "%s%d" % (name, chunk_id) in group.attrs:
																					data.extend(
																					    [n.decode('''utf8'''   ) if hasattr(SCREAMING_SNAKE_CASE    ,						'''decode'''   ) else n for n in group.attrs['''%s%d''' % (name, chunk_id)]]   )
																					chunk_id += 1
							return data



def 				lowerCAmelCase_			(      SCREAMING_SNAKE_CASE       :							Dict   ):


							'''simple docstring'''







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

							return tf.nest.map_structure(_expand_single_ad_tensor    ,						SCREAMING_SNAKE_CASE   )




