




'''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   UpperCamelCase_(			snake_case							: Dict=2							,       snake_case							: List[str]=True							,       snake_case							: Optional[Any]=False  ):

							'''simple docstring'''






							snake_case_			      =   Accelerator(step_scheduler_with_optimizer=snake_case							,       split_batches=snake_case  )
							snake_case_			      =   torch.nn.Linear(2							,       4  )
							snake_case_			      =   torch.optim.AdamW(model.parameters()							,       lr=1.0  )
							snake_case_			      =   torch.optim.lr_scheduler.OneCycleLR(snake_case							,       max_lr=0.01							,       steps_per_epoch=2							,       epochs=1  )
							snake_case_      ,   snake_case_      ,   snake_case_			      =   accelerator.prepare(snake_case							,       snake_case							,       snake_case  )

							# 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   UpperCamelCase_(			snake_case							: Any=2							,       snake_case							: List[Any]=True							,       snake_case							: str=False  ):

							'''simple docstring'''






							snake_case_			      =   Accelerator(step_scheduler_with_optimizer=snake_case							,       split_batches=snake_case  )
							snake_case_			      =   torch.nn.Linear(2							,       4  )
							snake_case_			      =   torch.optim.AdamW(model.parameters()							,       lr=1.0  )
							snake_case_			      =   torch.optim.lr_scheduler.LambdaLR(snake_case							,       lr_lambda=lambda snake_case  : 1 - n / 1_0  )
							snake_case_      ,   snake_case_      ,   snake_case_			      =   accelerator.prepare(snake_case							,       snake_case							,       snake_case  )

							# Optimizer has stepped
							snake_case_			      =   False
							scheduler.step()
							snake_case_			      =   1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 1_0
							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
							snake_case_			      =   True
							scheduler.step()
							if not step_scheduler_with_optimizer:
														snake_case_			      =   1 - 2 / 1_0
							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   UpperCamelCase_(			snake_case							: int = 2  ):

							'''simple docstring'''






							from transformers import get_linear_schedule_with_warmup

							snake_case_			      =   [1, 2, 4]
							for num_steps in steps:
														snake_case_			      =   GradientAccumulationPlugin(num_steps=snake_case							,       adjust_scheduler=num_steps > 1  )
														snake_case_			      =   Accelerator(gradient_accumulation_plugin=snake_case  )
														snake_case_			      =   torch.nn.Linear(2							,       4  )
														snake_case_			      =   torch.optim.AdamW(model.parameters()							,       lr=10.0  )
														snake_case_			      =   get_linear_schedule_with_warmup(optimizer=snake_case							,       num_warmup_steps=0							,       num_training_steps=2_0  )

														snake_case_      ,   snake_case_      ,   snake_case_			      =   accelerator.prepare(snake_case							,       snake_case							,       snake_case  )

														for i in range(1_0 * num_steps  ):
																					with accelerator.accumulate(snake_case  ):
																												optimizer.step()
																												scheduler.step()

																					if i == (1_0 * 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 				_snake_case	(			unittest.TestCase							):




			def      lowerCAmelCase__       (      self							)       ->       str:

										'''simple docstring'''




										debug_launcher(partial(a__       ,							num_processes=1							)       ,							num_processes=1							)
										debug_launcher(partial(a__       ,							num_processes=1       ,							split_batches=a__							)       ,							num_processes=1							)




			def      lowerCAmelCase__       (      self							)       ->       int:

										'''simple docstring'''




										debug_launcher(partial(a__       ,							num_processes=1							)       ,							num_processes=1							)
										debug_launcher(partial(a__       ,							num_processes=1       ,							split_batches=a__							)       ,							num_processes=1							)




			def      lowerCAmelCase__       (      self							)       ->       Optional[Any]:

										'''simple docstring'''




										debug_launcher(partial(a__       ,							num_processes=1       ,							step_scheduler_with_optimizer=a__							)       ,							num_processes=1							)




			def      lowerCAmelCase__       (      self							)       ->       int:

										'''simple docstring'''




										debug_launcher(partial(a__       ,							num_processes=1       ,							step_scheduler_with_optimizer=a__							)       ,							num_processes=1							)




			def      lowerCAmelCase__       (      self							)       ->       Optional[Any]:

										'''simple docstring'''




										AcceleratorState._reset_state(a__							)
										debug_launcher(a__							)
										debug_launcher(partial(a__       ,							num_processes=1       ,							split_batches=a__							)       ,							num_processes=1							)




			def      lowerCAmelCase__       (      self							)       ->       str:

										'''simple docstring'''




										AcceleratorState._reset_state(a__							)
										debug_launcher(a__							)
										debug_launcher(partial(a__       ,							num_processes=1       ,							split_batches=a__							)       ,							num_processes=1							)




			def      lowerCAmelCase__       (      self							)       ->       List[Any]:

										'''simple docstring'''




										AcceleratorState._reset_state(a__							)
										debug_launcher(partial(a__       ,							step_scheduler_with_optimizer=a__							)							)




			def      lowerCAmelCase__       (      self							)       ->       List[Any]:

										'''simple docstring'''




										AcceleratorState._reset_state(a__							)
										debug_launcher(partial(a__       ,							step_scheduler_with_optimizer=a__							)							)




			@require_huggingface_suite
			def      lowerCAmelCase__       (      self							)       ->       List[Any]:

										'''simple docstring'''




										AcceleratorState._reset_state(a__							)
										debug_launcher(partial(a__       ,							num_processes=1							)							)
										debug_launcher(a__							)










'''simple docstring'''





import random
from typing import List, Optional, Union

import numpy as np
import torch

from ..state import AcceleratorState
from .constants import CUDA_DISTRIBUTED_TYPES
from .dataclasses import DistributedType, RNGType
from .imports import is_npu_available, is_tpu_available, is_xpu_available


if is_tpu_available(check_device=False):
		import torch_xla.core.xla_model as xm


def   UpperCamelCase_(			snake_case							: int							,       snake_case							: bool = False  ):

							'''simple docstring'''






							if device_specific:
														seed += AcceleratorState().process_index
							random.seed(snake_case  )
							np.random.seed(snake_case  )
							torch.manual_seed(snake_case  )
							if is_xpu_available():
														torch.xpu.manual_seed_all(snake_case  )
							elif is_npu_available():
														torch.npu.manual_seed_all(snake_case  )
							else:
														torch.cuda.manual_seed_all(snake_case  )
							# ^^ safe to call this function even if cuda is not available
							if is_tpu_available():
														xm.set_rng_state(snake_case  )


def   UpperCamelCase_(			snake_case							: Optional[RNGType] = None							,       snake_case							: Optional[torch.Generator] = None  ):

							'''simple docstring'''






							if rng_type == RNGType.TORCH:
														snake_case_			      =   torch.get_rng_state()
							elif rng_type == RNGType.CUDA:
														snake_case_			      =   torch.cuda.get_rng_state()
							elif rng_type == RNGType.XLA:
														assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
														snake_case_			      =   torch.tensor(xm.get_rng_state()  )
							elif rng_type == RNGType.NPU:
														assert is_npu_available(), "Can't synchronize NPU seeds on an environment without NPUs."
														snake_case_			      =   torch.npu.get_rng_state()
							elif rng_type == RNGType.XPU:
														assert is_xpu_available(), "Can't synchronize XPU seeds on an environment without XPUs."
														snake_case_			      =   torch.xpu.get_rng_state()
							elif rng_type == RNGType.GENERATOR:
														assert generator is not None, "Need a generator to synchronize its seed."
														snake_case_			      =   generator.get_state()

							# Broadcast the rng state from device 0 to other devices
							snake_case_			      =   AcceleratorState()
							if state.distributed_type == DistributedType.TPU:
														snake_case_			      =   rng_state.to(xm.xla_device()  )
														xm.collective_broadcast([rng_state]  )
														xm.mark_step()
														snake_case_			      =   rng_state.cpu()
							elif (
							    state.distributed_type in CUDA_DISTRIBUTED_TYPES
							    or state.distributed_type == DistributedType.MULTI_NPU
							    or state.distributed_type == DistributedType.MULTI_XPU
							):
														snake_case_			      =   rng_state.to(state.device  )
														torch.distributed.broadcast(snake_case							,       0  )
														snake_case_			      =   rng_state.cpu()
							elif state.distributed_type == DistributedType.MULTI_CPU:
														torch.distributed.broadcast(snake_case							,       0  )

							# Set the broadcast rng state
							if rng_type == RNGType.TORCH:
														torch.set_rng_state(snake_case  )
							elif rng_type == RNGType.CUDA:
														torch.cuda.set_rng_state(snake_case  )
							elif rng_type == RNGType.NPU:
														torch.npu.set_rng_state(snake_case  )
							elif rng_type == RNGType.XPU:
														torch.xpu.set_rng_state(snake_case  )
							elif rng_type == RNGType.XLA:
														xm.set_rng_state(rng_state.item()  )
							elif rng_type == RNGType.GENERATOR:
														generator.set_state(snake_case  )




def   UpperCamelCase_(			snake_case							: List[Union[str, RNGType]]							,       snake_case							: Optional[torch.Generator] = None  ):

							'''simple docstring'''






							for rng_type in rng_types:
														synchronize_rng_state(RNGType(snake_case  )							,       generator=snake_case  )




