




import unittest

import numpy as np

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

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
		import torch

if is_vision_available():
		from PIL import Image

		from transformers import SwinaSRImageProcessor
		from transformers.image_transforms import get_image_size
class 		UpperCAmelCase       (    unittest.TestCase      ):







						'''simple docstring'''





						def __init__(						self							:					str	,     __lowercase							:					int	,     __lowercase							:					Tuple=7	,     __lowercase							:					List[Any]=3	,     __lowercase							:					List[str]=18	,     __lowercase							:					str=30	,     __lowercase							:					Dict=4_00	,     __lowercase							:					Optional[Any]=True	,     __lowercase							:					Any=1 / 2_55	,     __lowercase							:					Union[str, Any]=True	,     __lowercase							:					List[Any]=8	,     ):






							"""simple docstring"""






							snake_case_		   =					parent
							snake_case_		   =					batch_size
							snake_case_		   =					num_channels
							snake_case_		   =					image_size
							snake_case_		   =					min_resolution
							snake_case_		   =					max_resolution
							snake_case_		   =					do_rescale
							snake_case_		   =					rescale_factor
							snake_case_		   =					do_pad
							snake_case_		   =					pad_size






						def 	snake_case__						(						self							:					Tuple						):






							"""simple docstring"""






							return {
							    "do_rescale": self.do_rescale,
							    "rescale_factor": self.rescale_factor,
							    "do_pad": self.do_pad,
							    "pad_size": self.pad_size,
							}






						def 	snake_case__						(						self							:					Optional[int]	,     __lowercase							:					int=False	,     __lowercase							:					Tuple=False	,     __lowercase							:					Optional[int]=False						):






							"""simple docstring"""






							assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

							if equal_resolution:
								snake_case_		   =					[]
								for i in range(self.batch_size						):
									image_inputs.append(
									    np.random.randint(
									        2_55	,     size=(self.num_channels, self.max_resolution, self.max_resolution)	,     dtype=np.uinta						)						)
							else:
								snake_case_		   =					[]
								for i in range(self.batch_size						):
									snake_case_		   =					np.random.choice(np.arange(self.min_resolution	,     self.max_resolution						)	,     2						)
									image_inputs.append(np.random.randint(2_55	,     size=(self.num_channels, width, height)	,     dtype=np.uinta						)						)

							if not numpify and not torchify:
								# PIL expects the channel dimension as last dimension
								snake_case_		   =					[Image.fromarray(np.moveaxis(UpperCamelCase__	,     0	,     -1						)						) for x in image_inputs]

							if torchify:
								snake_case_		   =					[torch.from_numpy(UpperCamelCase__						) for x in image_inputs]

							return image_inputs







@require_torch
@require_vision
class 		UpperCAmelCase       (    lowerCAmelCase__	,       unittest.TestCase      ):







						'''simple docstring'''





						lowerCAmelCase_							=       SwinaSRImageProcessor if is_vision_available() else None






						def 	snake_case__						(						self							:					Optional[int]						):






							"""simple docstring"""






							snake_case_		   =					SwinaSRImageProcessingTester(self						)






						@property
						def 	snake_case__						(						self							:					List[str]						):






							"""simple docstring"""






							return self.image_processor_tester.prepare_image_processor_dict()






						def 	snake_case__						(						self							:					List[Any]						):






							"""simple docstring"""






							snake_case_		   =					self.image_processing_class(**self.image_processor_dict						)
							self.assertTrue(hasattr(UpperCamelCase__	,     "do_rescale"						)						)
							self.assertTrue(hasattr(UpperCamelCase__	,     "rescale_factor"						)						)
							self.assertTrue(hasattr(UpperCamelCase__	,     "do_pad"						)						)
							self.assertTrue(hasattr(UpperCamelCase__	,     "pad_size"						)						)






						def 	snake_case__						(						self							:					List[Any]						):






							"""simple docstring"""






							pass






						def 	snake_case__						(						self							:					str	,     __lowercase							:					Any						):






							"""simple docstring"""






							snake_case_		   =					get_image_size(UpperCamelCase__						)
							snake_case_		   =					self.image_processor_tester.pad_size

							snake_case_		   =					(old_height // size + 1) * size - old_height
							snake_case_		   =					(old_width // size + 1) * size - old_width
							return old_height + pad_height, old_width + pad_width






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






							"""simple docstring"""






							snake_case_		   =					self.image_processing_class(**self.image_processor_dict						)
							# create random PIL images
							snake_case_		   =					self.image_processor_tester.prepare_inputs(equal_resolution=UpperCamelCase__						)
							for image in image_inputs:
								self.assertIsInstance(UpperCamelCase__	,     Image.Image						)

							# Test not batched input
							snake_case_		   =					image_processor(image_inputs[0]	,     return_tensors="pt"						).pixel_values
							snake_case_		   =					self.calculate_expected_size(np.array(image_inputs[0]						)						)
							self.assertEqual(
							    encoded_images.shape	,     (
							        1,
							        self.image_processor_tester.num_channels,
							        expected_height,
							        expected_width,
							    )	,     )






						def 	snake_case__						(						self							:					str						):






							"""simple docstring"""






							snake_case_		   =					self.image_processing_class(**self.image_processor_dict						)
							# create random numpy tensors
							snake_case_		   =					self.image_processor_tester.prepare_inputs(equal_resolution=UpperCamelCase__	,     numpify=UpperCamelCase__						)
							for image in image_inputs:
								self.assertIsInstance(UpperCamelCase__	,     np.ndarray						)

							# Test not batched input
							snake_case_		   =					image_processor(image_inputs[0]	,     return_tensors="pt"						).pixel_values
							snake_case_		   =					self.calculate_expected_size(image_inputs[0]						)
							self.assertEqual(
							    encoded_images.shape	,     (
							        1,
							        self.image_processor_tester.num_channels,
							        expected_height,
							        expected_width,
							    )	,     )






						def 	snake_case__						(						self							:					Tuple						):






							"""simple docstring"""






							snake_case_		   =					self.image_processing_class(**self.image_processor_dict						)
							# create random PyTorch tensors
							snake_case_		   =					self.image_processor_tester.prepare_inputs(equal_resolution=UpperCamelCase__	,     torchify=UpperCamelCase__						)
							for image in image_inputs:
								self.assertIsInstance(UpperCamelCase__	,     torch.Tensor						)

							# Test not batched input
							snake_case_		   =					image_processor(image_inputs[0]	,     return_tensors="pt"						).pixel_values
							snake_case_		   =					self.calculate_expected_size(image_inputs[0]						)
							self.assertEqual(
							    encoded_images.shape	,     (
							        1,
							        self.image_processor_tester.num_channels,
							        expected_height,
							        expected_width,
							    )	,     )






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 						lowerCamelCase__   (		_A=2					,     _A=True					,     _A=False							):







	'''simple docstring'''





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

	# 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 						lowerCamelCase__   (		_A=2					,     _A=True					,     _A=False							):







	'''simple docstring'''





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

	# Optimizer has stepped
	snake_case_		   =					False
	scheduler.step()
	snake_case_		   =					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
	snake_case_		   =					True
	scheduler.step()
	if not step_scheduler_with_optimizer:
		snake_case_		   =					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 						lowerCamelCase__   (		_A = 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=_A					,     adjust_scheduler=num_steps > 1							)
		snake_case_		   =					Accelerator(gradient_accumulation_plugin=_A							)
		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=_A					,     num_warmup_steps=0					,     num_training_steps=20							)

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

		for i in range(10 * num_steps							):
			with accelerator.accumulate(_A							):
				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 		UpperCAmelCase       (    unittest.TestCase      ):







						'''simple docstring'''





						def 	snake_case__						(						self							:					Tuple						):






							"""simple docstring"""






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






						def 	snake_case__						(						self							:					int						):






							"""simple docstring"""






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






						def 	snake_case__						(						self							:					Dict						):






							"""simple docstring"""






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






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






							"""simple docstring"""






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






						def 	snake_case__						(						self							:					List[Any]						):






							"""simple docstring"""






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






						def 	snake_case__						(						self							:					Dict						):






							"""simple docstring"""






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






						def 	snake_case__						(						self							:					List[str]						):






							"""simple docstring"""






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






						def 	snake_case__						(						self							:					Optional[Any]						):






							"""simple docstring"""






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






						@require_huggingface_suite
						def 	snake_case__						(						self							:					Optional[Any]						):






							"""simple docstring"""






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