



'''simple docstring'''






import unittest
from functools import partial

import torch

from accelerate import Accelerator, debug_launcher
from accelerate.state import AcceleratorState, GradientState
from accelerate.test_utils import require_cpu, require_huggingface_suite
from accelerate.utils import GradientAccumulationPlugin




def        snake_case_   (	lowerCAmelCase_=2			,  lowerCAmelCase_=True			,  lowerCAmelCase_=False						)->						Any:





       '''simple docstring'''


       _UpperCAmelCase    :		List[str]     	=					Accelerator(step_scheduler_with_optimizer=lowerCAmelCase_			,  split_batches=lowerCAmelCase_						)
       _UpperCAmelCase    :		Union[str, Any]     	=					torch.nn.Linear(2			,  4						)
       _UpperCAmelCase    :		Tuple     	=					torch.optim.AdamW(model.parameters()			,  lr=1.0						)
       _UpperCAmelCase    :		List[Any]     	=					torch.optim.lr_scheduler.OneCycleLR(lowerCAmelCase_			,  max_lr=0.0_1			,  steps_per_epoch=2			,  epochs=1						)
       _UpperCAmelCase	,_UpperCAmelCase	,_UpperCAmelCase    :		str     	=					accelerator.prepare(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_						)

       # Optimizer has stepped
       scheduler.step()
       if step_scheduler_with_optimizer or (num_processes == 1):
              assert (
                  scheduler.scheduler.last_epoch == num_processes
              ), F'''Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})'''
       else:
              assert (
                  scheduler.scheduler.last_epoch != num_processes
              ), F'''Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})'''




def        snake_case_   (	lowerCAmelCase_=2			,  lowerCAmelCase_=True			,  lowerCAmelCase_=False						)->						List[Any]:





       '''simple docstring'''


       _UpperCAmelCase    :		List[Any]     	=					Accelerator(step_scheduler_with_optimizer=lowerCAmelCase_			,  split_batches=lowerCAmelCase_						)
       _UpperCAmelCase    :		str     	=					torch.nn.Linear(2			,  4						)
       _UpperCAmelCase    :		List[str]     	=					torch.optim.AdamW(model.parameters()			,  lr=1.0						)
       _UpperCAmelCase    :		Union[str, Any]     	=					torch.optim.lr_scheduler.LambdaLR(lowerCAmelCase_			,  lr_lambda=lambda lowerCAmelCase_						: 1 - n / 10						)
       _UpperCAmelCase	,_UpperCAmelCase	,_UpperCAmelCase    :		int     	=					accelerator.prepare(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_						)

       # Optimizer has stepped
       _UpperCAmelCase    :		int     	=					False
       scheduler.step()
       _UpperCAmelCase    :		int     	=					1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
       assert (
           scheduler.get_last_lr()[0] == expected_lr
       ), F'''Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'''

       # Optimizer has not stepped
       _UpperCAmelCase    :		Any     	=					True
       scheduler.step()
       if not step_scheduler_with_optimizer:
              _UpperCAmelCase    :		Optional[int]     	=					1 - 2 / 10
       assert (
           scheduler.get_last_lr()[0] == expected_lr
       ), F'''Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'''




def        snake_case_   (	lowerCAmelCase_ = 2						)->						Tuple:





       '''simple docstring'''


       from transformers import get_linear_schedule_with_warmup

       _UpperCAmelCase    :		Dict     	=					[1, 2, 4]
       for num_steps in steps:
              _UpperCAmelCase    :		Union[str, Any]     	=					GradientAccumulationPlugin(num_steps=lowerCAmelCase_			,  adjust_scheduler=num_steps > 1						)
              _UpperCAmelCase    :		Optional[int]     	=					Accelerator(gradient_accumulation_plugin=lowerCAmelCase_						)
              _UpperCAmelCase    :		Any     	=					torch.nn.Linear(2			,  4						)
              _UpperCAmelCase    :		Optional[Any]     	=					torch.optim.AdamW(model.parameters()			,  lr=1_0.0						)
              _UpperCAmelCase    :		Optional[int]     	=					get_linear_schedule_with_warmup(optimizer=lowerCAmelCase_			,  num_warmup_steps=0			,  num_training_steps=20						)

              _UpperCAmelCase	,_UpperCAmelCase	,_UpperCAmelCase    :		Tuple     	=					accelerator.prepare(lowerCAmelCase_			,  lowerCAmelCase_			,  lowerCAmelCase_						)

              for i in range(10 * num_steps						):
                     with accelerator.accumulate(lowerCAmelCase_						):
                            optimizer.step()
                            scheduler.step()

                     if i == (10 * num_steps - 2):
                            assert (
                                scheduler.get_last_lr()[0] != 0
                            ), F'''Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}'''
              assert (
                  scheduler.get_last_lr()[0] == 0
              ), F'''Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}'''
              GradientState._reset_state()



@require_cpu
class 				lowercase			(    unittest.TestCase							):


 """simple docstring"""


 def _snake_case					(				self				)						->  Union[str, Any]:
        debug_launcher(partial(a_    ,num_processes=1				)    ,num_processes=1				)
        debug_launcher(partial(a_    ,num_processes=1    ,split_batches=a_				)    ,num_processes=1				)







 def _snake_case					(				self				)						->  Tuple:
        debug_launcher(partial(a_    ,num_processes=1				)    ,num_processes=1				)
        debug_launcher(partial(a_    ,num_processes=1    ,split_batches=a_				)    ,num_processes=1				)







 def _snake_case					(				self				)						->  List[Any]:
        debug_launcher(partial(a_    ,num_processes=1    ,step_scheduler_with_optimizer=a_				)    ,num_processes=1				)







 def _snake_case					(				self				)						->  Optional[int]:
        debug_launcher(partial(a_    ,num_processes=1    ,step_scheduler_with_optimizer=a_				)    ,num_processes=1				)







 def _snake_case					(				self				)						->  List[Any]:
        AcceleratorState._reset_state(a_				)
        debug_launcher(a_				)
        debug_launcher(partial(a_    ,num_processes=1    ,split_batches=a_				)    ,num_processes=1				)







 def _snake_case					(				self				)						->  int:
        AcceleratorState._reset_state(a_				)
        debug_launcher(a_				)
        debug_launcher(partial(a_    ,num_processes=1    ,split_batches=a_				)    ,num_processes=1				)







 def _snake_case					(				self				)						->  Optional[int]:
        AcceleratorState._reset_state(a_				)
        debug_launcher(partial(a_    ,step_scheduler_with_optimizer=a_				)				)







 def _snake_case					(				self				)						->  Dict:
        AcceleratorState._reset_state(a_				)
        debug_launcher(partial(a_    ,step_scheduler_with_optimizer=a_				)				)







 @require_huggingface_suite
 def _snake_case					(				self				)						->  Any:
        AcceleratorState._reset_state(a_				)
        debug_launcher(partial(a_    ,num_processes=1				)				)
        debug_launcher(a_				)








'''simple docstring'''






from unittest.mock import patch

import datasets
from datasets import Dataset




def        snake_case_   (	)->						str:





       '''simple docstring'''


       _UpperCAmelCase    :		Tuple     	=					Dataset.from_dict({"""col_1""": [3, 2, 0, 1]}						)

       with patch("""tqdm.auto.tqdm"""						) as mock_tqdm:
              datasets.disable_progress_bar()
              dset.map(lambda lowerCAmelCase_						: {"col_2": x["col_1"] + 1}						)
              mock_tqdm.assert_not_called()

              mock_tqdm.reset_mock()

              datasets.enable_progress_bar()
              dset.map(lambda lowerCAmelCase_						: {"col_2": x["col_1"] + 1}						)
              mock_tqdm.assert_called()



