



"""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 	snake_case_  (				A_ :		str,    A_ :		List[str]=10						):


       '''simple docstring'''





       _lowerCamelCase						:    Optional[int]         =			[]
       for _ in range(a_						):
              lrs.append(scheduler.get_lr()[0]						)
              scheduler.step()
       return lrs






def 	snake_case_  (				A_ :		int,    A_ :		int=10						):


       '''simple docstring'''





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

                            _lowerCamelCase						:    Union[str, Any]         =			torch.load(a_						)
                            scheduler.load_state_dict(a_						)
       return lrs






@require_torch
class 	__snake_case	(     unittest.TestCase):




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						int		):






         """simple docstring"""






         self.assertEqual(len(lowerCAmelCase__		)						, len(lowerCAmelCase__		)		)
         for a, b in zip(lowerCAmelCase__						, lowerCAmelCase__		):
                self.assertAlmostEqual(lowerCAmelCase__						, lowerCAmelCase__						, delta=lowerCAmelCase__		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						str		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			torch.tensor([0.1, -0.2, -0.1]						, requires_grad=lowerCAmelCase__		)
         _lowerCamelCase						:    Optional[int]         =			torch.tensor([0.4, 0.2, -0.5]		)
         _lowerCamelCase						:    List[Any]         =			nn.MSELoss()
         # No warmup, constant schedule, no gradient clipping
         _lowerCamelCase						:    str         =			AdamW(params=[w]						, lr=2E-1						, weight_decay=0.0		)
         for _ in range(1_0_0		):
                _lowerCamelCase						:    Tuple         =			criterion(lowerCAmelCase__						, lowerCAmelCase__		)
                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 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			torch.tensor([0.1, -0.2, -0.1]						, requires_grad=lowerCAmelCase__		)
         _lowerCamelCase						:    Tuple         =			torch.tensor([0.4, 0.2, -0.5]		)
         _lowerCamelCase						:    Dict         =			nn.MSELoss()
         # No warmup, constant schedule, no gradient clipping
         _lowerCamelCase						:    Any         =			Adafactor(
             params=[w]						, lr=1E-2						, eps=(1E-30, 1E-3)						, clip_threshold=1.0						, decay_rate=-0.8						, betaa=lowerCAmelCase__						, weight_decay=0.0						, relative_step=lowerCAmelCase__						, scale_parameter=lowerCAmelCase__						, warmup_init=lowerCAmelCase__						, )
         for _ in range(1_0_0_0		):
                _lowerCamelCase						:    Optional[int]         =			criterion(lowerCAmelCase__						, lowerCAmelCase__		)
                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 	__snake_case	(     unittest.TestCase):
  snake_case__    :       Union[str, Any]        =						nn.Linear(5_0  ,					5_0) if is_torch_available() else None
  snake_case__    :       Dict        =						AdamW(m.parameters()  ,					lr=10.0) if is_torch_available() else None
  snake_case__    :       Optional[int]        =						1_0




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						List[Any]=None		):






         """simple docstring"""






         self.assertEqual(len(lowerCAmelCase__		)						, len(lowerCAmelCase__		)		)
         for a, b in zip(lowerCAmelCase__						, lowerCAmelCase__		):
                self.assertAlmostEqual(lowerCAmelCase__						, lowerCAmelCase__						, delta=lowerCAmelCase__						, msg=lowerCAmelCase__		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			{"num_warmup_steps": 2, "num_training_steps": 1_0}
         # 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.6_56, 5.6_25, 3.9_06, 2.5, 1.4_06, 0.6_25, 0.1_56],
             ),
             get_inverse_sqrt_schedule: (
                 {"num_warmup_steps": 2},
                 [0.0, 5.0, 10.0, 8.1_65, 7.0_71, 6.3_25, 5.7_74, 5.3_45, 5.0, 4.7_14],
             ),
         }

         for scheduler_func, data in scheds.items():
                _lowerCamelCase						:    Dict         =			data

                _lowerCamelCase						:    Optional[Any]         =			scheduler_func(self.optimizer						, **lowerCAmelCase__		)
                self.assertEqual(len([scheduler.get_lr()[0]]		)						, 1		)
                _lowerCamelCase						:    Optional[Any]         =			unwrap_schedule(lowerCAmelCase__						, self.num_steps		)
                self.assertListAlmostEqual(
                    lowerCAmelCase__						, lowerCAmelCase__						, tol=1E-2						, msg=f'''failed for {scheduler_func} in normal scheduler'''						, )

                _lowerCamelCase						:    Any         =			scheduler_func(self.optimizer						, **lowerCAmelCase__		)
                if scheduler_func.__name__ != "get_constant_schedule":
                       LambdaScheduleWrapper.wrap_scheduler(lowerCAmelCase__		)  # wrap to test picklability of the schedule
                _lowerCamelCase						:    Optional[Any]         =			unwrap_and_save_reload_schedule(lowerCAmelCase__						, self.num_steps		)
                self.assertListEqual(lowerCAmelCase__						, lowerCAmelCase__						, msg=f'''failed for {scheduler_func} in save and reload'''		)






class 	__snake_case	:




  def __init__(			self					:						Dict						, __lowerCAmelCase					:						Union[str, Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			fn




  def __call__(			self					:						Tuple						, *__lowerCAmelCase					:						int						, **__lowerCAmelCase					:						Tuple		):






         """simple docstring"""






         return self.fn(*lowerCAmelCase__						, **lowerCAmelCase__		)




  @classmethod
  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						Dict		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			list(map(self						, scheduler.lr_lambdas		)		)







"""simple docstring"""






import argparse


lowerCAmelCase__      =    '''docs/source/_static/js/custom.js'''






def 	snake_case_  (				A_ :		List[str]						):


       '''simple docstring'''





       with open(A_,    encoding='''utf-8''',    newline='''\n'''						) as f:
              _lowerCamelCase						:    int         =			f.readlines()
       _lowerCamelCase						:    List[str]         =			0

       # First let's put the right version
       while not lines[index].startswith('''const stableVersion ='''						):
              index += 1
       _lowerCamelCase						:    List[Any]         =			F'''const stableVersion = "v{version}"\n'''

       # Then update the dictionary
       while not lines[index].startswith('''const versionMapping = {'''						):
              index += 1

       # We go until the end
       while not lines[index].startswith('''}'''						):
              index += 1
       # We add the new version at the end
       lines[index - 1] += F'''    "v{version}": "v{version}",\n'''

       with open(A_,    '''w''',    encoding='''utf-8''',    newline='''\n'''						) as f:
              f.writelines(A_						)


if __name__ == "__main__":
   lowerCAmelCase__      =    argparse.ArgumentParser()
   parser.add_argument('''--version''', help='''Release version.''')
   lowerCAmelCase__      =    parser.parse_args()
   update_custom_js(args.version)


