





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        __SCREAMING_SNAKE_CASE	(				__lowercase:		Any=2  ,				__lowercase:		Optional[int]=True  ,				__lowercase:		Any=False):



  '''simple docstring'''






  __lowerCAmelCase								=					Accelerator(step_scheduler_with_optimizer=__lowerCAmelCase  ,				split_batches=__lowerCAmelCase)
  __lowerCAmelCase								=					torch.nn.Linear(2  ,				4)
  __lowerCAmelCase								=					torch.optim.AdamW(model.parameters()  ,				lr=1.0)
  __lowerCAmelCase								=					torch.optim.lr_scheduler.OneCycleLR(__lowerCAmelCase  ,				max_lr=0.01  ,				steps_per_epoch=2  ,				epochs=1)
  __lowerCAmelCase								=					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        __SCREAMING_SNAKE_CASE	(				__lowercase:		Optional[int]=2  ,				__lowercase:		Union[str, Any]=True  ,				__lowercase:		int=False):



  '''simple docstring'''






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

  # Optimizer has stepped
  __lowerCAmelCase								=					False
  scheduler.step()
  __lowerCAmelCase								=					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
  __lowerCAmelCase								=					True
  scheduler.step()
  if not step_scheduler_with_optimizer:
    __lowerCAmelCase								=					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        __SCREAMING_SNAKE_CASE	(				__lowercase:		int = 2):



  '''simple docstring'''






  from transformers import get_linear_schedule_with_warmup

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

    __lowerCAmelCase								=					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 	__snake_case					(      unittest.TestCase ):






    '''simple docstring'''





    def 						A		( self    :		Any				):

      """simple docstring"""





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





    def 						A		( self    :		Union[str, Any]				):

      """simple docstring"""





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





    def 						A		( self    :		Dict				):

      """simple docstring"""





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





    def 						A		( self    :		str				):

      """simple docstring"""





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





    def 						A		( self    :		Optional[int]				):

      """simple docstring"""





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





    def 						A		( self    :		Dict				):

      """simple docstring"""





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





    def 						A		( self    :		int				):

      """simple docstring"""





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





    def 						A		( self    :		Optional[int]				):

      """simple docstring"""





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





    @require_huggingface_suite
    def 						A		( self    :		Tuple				):

      """simple docstring"""





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






"""simple docstring"""
import json
import os
import tempfile

from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available


if is_torch_available():
    import numpy as np
    import torch

if is_vision_available():
    from PIL import Image






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Any  ,	__lowerCAmelCase			:		Dict=False  ,	__lowerCAmelCase			:		Optional[Any]=False  ,	__lowerCAmelCase			:		int=False						)    ->       List[str]:
 assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

 _UpperCamelCase      :	Optional[int]					      =  []
 for i in range(image_processor_tester.batch_size						):
  if equal_resolution:
   _UpperCamelCase      :	List[str]					      =  image_processor_tester.max_resolution
  else:
   # To avoid getting image width/height 0
   _UpperCamelCase      :	int					      =  image_processor_tester.min_resolution
   if getattr(__lowerCAmelCase  ,	"size_divisor"  ,	__lowerCAmelCase						):
    # If `size_divisor` is defined, the image needs to have width/size >= `size_divisor`
    _UpperCamelCase      :	Optional[int]					      =  max(image_processor_tester.size_divisor  ,	__lowerCAmelCase						)
   _UpperCamelCase      ,	_UpperCamelCase      :	Optional[Any]					      =  np.random.choice(np.arange(__lowerCAmelCase  ,	image_processor_tester.max_resolution						)  ,	2						)
  image_inputs.append(
      np.random.randint(255  ,	size=(image_processor_tester.num_channels, width, height)  ,	dtype=np.uinta						)						)

 if not numpify and not torchify:
  # PIL expects the channel dimension as last dimension
  _UpperCamelCase      :	Optional[Any]					      =  [Image.fromarray(np.moveaxis(__lowerCAmelCase  ,	0  ,	-1						)						) for image in image_inputs]

 if torchify:
  _UpperCamelCase      :	str					      =  [torch.from_numpy(__lowerCAmelCase						) for image in image_inputs]

 return image_inputs






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		Optional[Any]  ,	__lowerCAmelCase			:		Tuple=10  ,	__lowerCAmelCase			:		str=10  ,	__lowerCAmelCase			:		Any=False  ,	__lowerCAmelCase			:		List[Any]=False						)    ->       Optional[int]:
 _UpperCamelCase      :	Optional[int]					      =  []
 for i in range(image_processor_tester.num_frames						):
  video.append(np.random.randint(255  ,	size=(image_processor_tester.num_channels, width, height)  ,	dtype=np.uinta						)						)

 if not numpify and not torchify:
  # PIL expects the channel dimension as last dimension
  _UpperCamelCase      :	Optional[int]					      =  [Image.fromarray(np.moveaxis(__lowerCAmelCase  ,	0  ,	-1						)						) for frame in video]

 if torchify:
  _UpperCamelCase      :	int					      =  [torch.from_numpy(__lowerCAmelCase						) for frame in video]

 return video






def 				__lowerCAmelCase					(					__lowerCAmelCase			:		List[Any]  ,	__lowerCAmelCase			:		Union[str, Any]=False  ,	__lowerCAmelCase			:		Any=False  ,	__lowerCAmelCase			:		Dict=False						)    ->       Any:
 assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

 _UpperCamelCase      :	Optional[Any]					      =  []
 for i in range(image_processor_tester.batch_size						):
  if equal_resolution:
   _UpperCamelCase      :	List[Any]					      =  image_processor_tester.max_resolution
  else:
   _UpperCamelCase      ,	_UpperCamelCase      :	Optional[int]					      =  np.random.choice(
       np.arange(image_processor_tester.min_resolution  ,	image_processor_tester.max_resolution						)  ,	2						)
   _UpperCamelCase      :	List[Any]					      =  prepare_video(
       image_processor_tester=__lowerCAmelCase  ,	width=__lowerCAmelCase  ,	height=__lowerCAmelCase  ,	numpify=__lowerCAmelCase  ,	torchify=__lowerCAmelCase  ,	)
  video_inputs.append(__lowerCAmelCase						)

 return video_inputs







class  _SCREAMING_SNAKE_CASE       :



 '''simple docstring'''





 __UpperCAmelCase		   =							None






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	str					      =  self.image_processing_class(**self.image_processor_dict   )
  _UpperCamelCase      :	Optional[int]					      =  json.loads(image_processor.to_json_string()   )
  for key, value in self.image_processor_dict.items():
   self.assertEqual(obj[key]					,     lowerCAmelCase__   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Union[str, Any]					      =  self.image_processing_class(**self.image_processor_dict   )

  with tempfile.TemporaryDirectory() as tmpdirname:
   _UpperCamelCase      :	Tuple					      =  os.path.join(lowerCAmelCase__					,     "image_processor.json"   )
   image_processor_first.to_json_file(lowerCAmelCase__   )
   _UpperCamelCase      :	Tuple					      =  self.image_processing_class.from_json_file(lowerCAmelCase__   )

  self.assertEqual(image_processor_second.to_dict()					,     image_processor_first.to_dict()   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Dict					      =  self.image_processing_class(**self.image_processor_dict   )

  with tempfile.TemporaryDirectory() as tmpdirname:
   _UpperCamelCase      :	Optional[int]					      =  image_processor_first.save_pretrained(lowerCAmelCase__   )[0]
   check_json_file_has_correct_format(lowerCAmelCase__   )
   _UpperCamelCase      :	int					      =  self.image_processing_class.from_pretrained(lowerCAmelCase__   )

  self.assertEqual(image_processor_second.to_dict()					,     image_processor_first.to_dict()   )






 def lowercase_      (self   ):

  '''simple docstring'''







  _UpperCamelCase      :	Tuple					      =  self.image_processing_class()
  self.assertIsNotNone(lowerCAmelCase__   )






 @require_torch
 @require_vision
 def lowercase_      (self   ):

  '''simple docstring'''







  if self.test_cast_dtype is not None:
   # Initialize image_processor
   _UpperCamelCase      :	Optional[int]					      =  self.image_processing_class(**self.image_processor_dict   )

   # create random PyTorch tensors
   _UpperCamelCase      :	str					      =  prepare_image_inputs(self.image_processor_tester					,     equal_resolution=lowerCAmelCase__					,     torchify=lowerCAmelCase__   )

   _UpperCamelCase      :	List[Any]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   )
   # for layoutLM compatiblity
   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.floataa   )

   _UpperCamelCase      :	List[Any]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   ).to(torch.floataa   )
   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.floataa   )

   _UpperCamelCase      :	str					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   ).to("cpu"					,     torch.bfloataa   )
   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.bfloataa   )

   with self.assertRaises(lowerCAmelCase__   ):
    _UpperCamelCase      :	Optional[int]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   ).to(torch.bfloataa					,     "cpu"   )

   # Try with text + image feature
   _UpperCamelCase      :	Union[str, Any]					      =  image_processor(lowerCAmelCase__					,     return_tensors="pt"   )
   encoding.update({"input_ids": torch.LongTensor([[1, 2, 3], [4, 5, 6]]   )}   )
   _UpperCamelCase      :	Optional[int]					      =  encoding.to(torch.floataa   )

   self.assertEqual(encoding.pixel_values.device					,     torch.device("cpu"   )   )
   self.assertEqual(encoding.pixel_values.dtype					,     torch.floataa   )
   self.assertEqual(encoding.input_ids.dtype					,     torch.long   )



