



import gc
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, KandinskyVaaPipeline, KandinskyVaaPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()

class   _A   ( _lowercase  ,						unittest.TestCase	):


						'''simple docstring'''





						_snake_case			:     str      							=      KandinskyVaaPipeline
						_snake_case			:     List[Any]      							=      [
						    """image_embeds""",
						    """negative_image_embeds""",
						]
						_snake_case			:     Optional[int]      							=      ["""image_embeds""", """negative_image_embeds"""]
						_snake_case			:     List[str]      							=      [
						    """generator""",
						    """height""",
						    """width""",
						    """latents""",
						    """guidance_scale""",
						    """num_inference_steps""",
						    """return_dict""",
						    """guidance_scale""",
						    """num_images_per_prompt""",
						    """output_type""",
						    """return_dict""",
						]
						_snake_case			:     str      							=      False



						@property
						def     _snake_case	(	self	:							Dict	):






							'''simple docstring'''



							return 32



						@property
						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							return 32



						@property
						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							return self.time_input_dim



						@property
						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							return self.time_input_dim * 4



						@property
						def     _snake_case	(	self	:							Optional[Any]	):






							'''simple docstring'''



							return 100



						@property
						def     _snake_case	(	self	:							Optional[Any]	):






							'''simple docstring'''



							torch.manual_seed(0	)

							__lowercase		       =			{
							    "in_channels": 4,
							    # Out channels is double in channels because predicts mean and variance
							    "out_channels": 8,
							    "addition_embed_type": "image",
							    "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
							    "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
							    "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
							    "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
							    "layers_per_block": 1,
							    "encoder_hid_dim": self.text_embedder_hidden_size,
							    "encoder_hid_dim_type": "image_proj",
							    "cross_attention_dim": self.cross_attention_dim,
							    "attention_head_dim": 4,
							    "resnet_time_scale_shift": "scale_shift",
							    "class_embed_type": None,
							}

							__lowercase		       =			UNetaDConditionModel(**lowerCamelCase	)
							return model



						@property
						def     _snake_case	(	self	:							Tuple	):






							'''simple docstring'''



							return {
							    "block_out_channels": [32, 64],
							    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
							    "in_channels": 3,
							    "latent_channels": 4,
							    "layers_per_block": 1,
							    "norm_num_groups": 8,
							    "norm_type": "spatial",
							    "num_vq_embeddings": 12,
							    "out_channels": 3,
							    "up_block_types": [
							        "AttnUpDecoderBlock2D",
							        "UpDecoderBlock2D",
							    ],
							    "vq_embed_dim": 4,
							}



						@property
						def     _snake_case	(	self	:							Dict	):






							'''simple docstring'''



							torch.manual_seed(0	)
							__lowercase		       =			VQModel(**self.dummy_movq_kwargs	)
							return model



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






							'''simple docstring'''



							__lowercase		       =			self.dummy_unet
							__lowercase		       =			self.dummy_movq

							__lowercase		       =			DDIMScheduler(
							    num_train_timesteps=1_000							,      beta_schedule="linear"							,      beta_start=0.0_0085							,      beta_end=0.012							,      clip_sample=lowerCamelCase							,      set_alpha_to_one=lowerCamelCase							,      steps_offset=1							,      prediction_type="epsilon"							,      thresholding=lowerCamelCase							,      )

							__lowercase		       =			{
							    "unet": unet,
							    "scheduler": scheduler,
							    "movq": movq,
							}
							return components



						def     _snake_case	(	self	:							Tuple							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							Tuple=0	):






							'''simple docstring'''



							__lowercase		       =			floats_tensor((1, self.text_embedder_hidden_size)							,      rng=random.Random(lowerCamelCase	)	).to(lowerCamelCase	)
							__lowercase		       =			floats_tensor((1, self.text_embedder_hidden_size)							,      rng=random.Random(seed + 1	)	).to(
							    lowerCamelCase	)
							if str(lowerCamelCase	).startswith("mps"	):
								__lowercase		       =			torch.manual_seed(lowerCamelCase	)
							else:
								__lowercase		       =			torch.Generator(device=lowerCamelCase	).manual_seed(lowerCamelCase	)
							__lowercase		       =			{
							    "image_embeds": image_embeds,
							    "negative_image_embeds": negative_image_embeds,
							    "generator": generator,
							    "height": 64,
							    "width": 64,
							    "guidance_scale": 4.0,
							    "num_inference_steps": 2,
							    "output_type": "np",
							}
							return inputs



						def     _snake_case	(	self	:							Optional[Any]	):






							'''simple docstring'''



							__lowercase		       =			"cpu"

							__lowercase		       =			self.get_dummy_components()

							__lowercase		       =			self.pipeline_class(**lowerCamelCase	)
							__lowercase		       =			pipe.to(lowerCamelCase	)

							pipe.set_progress_bar_config(disable=lowerCamelCase	)

							__lowercase		       =			pipe(**self.get_dummy_inputs(lowerCamelCase	)	)
							__lowercase		       =			output.images

							__lowercase		       =			pipe(
							    **self.get_dummy_inputs(lowerCamelCase	)							,      return_dict=lowerCamelCase							,      )[0]

							__lowercase		       =			image[0, -3:, -3:, -1]
							__lowercase		       =			image_from_tuple[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)

							__lowercase		       =			np.array(
							    [0.623_7976, 1.0, 0.3644_1332, 1.0, 0.7063_9634, 0.2987_7186, 0.8565_2125, 0.521_6843, 0.5445_4046]	)

							assert (
							    np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
							), f""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""

							assert (
							    np.abs(image_from_tuple_slice.flatten() - expected_slice	).max() < 1e-2
							), f""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""

@slow
@require_torch_gpu
class   _A   ( unittest.TestCase	):


						'''simple docstring'''





						def     _snake_case	(	self	:							str	):






							'''simple docstring'''



							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



						def     _snake_case	(	self	:							int	):






							'''simple docstring'''



							__lowercase		       =			load_numpy(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
							    "/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy"	)

							__lowercase		       =			KandinskyVaaPriorPipeline.from_pretrained(
							    "kandinsky-community/kandinsky-2-2-prior"							,      torch_dtype=torch.floataa	)
							pipe_prior.to(lowerCamelCase	)

							__lowercase		       =			KandinskyVaaPipeline.from_pretrained(
							    "kandinsky-community/kandinsky-2-2-decoder"							,      torch_dtype=torch.floataa	)
							__lowercase		       =			pipeline.to(lowerCamelCase	)
							pipeline.set_progress_bar_config(disable=lowerCamelCase	)

							__lowercase		       =			"red cat, 4k photo"

							__lowercase		       =			torch.Generator(device="cuda"	).manual_seed(0	)
							__lowercase       ,    __lowercase		       =			pipe_prior(
							    lowerCamelCase							,      generator=lowerCamelCase							,      num_inference_steps=5							,      negative_prompt=""							,      ).to_tuple()

							__lowercase		       =			torch.Generator(device="cuda"	).manual_seed(0	)
							__lowercase		       =			pipeline(
							    image_embeds=lowerCamelCase							,      negative_image_embeds=lowerCamelCase							,      generator=lowerCamelCase							,      num_inference_steps=100							,      output_type="np"							,      )

							__lowercase		       =			output.images[0]

							assert image.shape == (512, 512, 3)

							assert_mean_pixel_difference(lowerCamelCase							,      lowerCamelCase	)








from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed



def    snake_case_				(    _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   ):
	for param, grad_param in zip(model_a.parameters()  ,	model_b.parameters()   ):
		if not param.requires_grad:
			continue
		if not did_step:
			# Grads should not be in sync
			assert (
			    torch.allclose(param.grad  ,	grad_param.grad   ) is False
			), F"""Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})"""
		else:
			# Grads should be in sync
			assert (
			    torch.allclose(param.grad  ,	grad_param.grad   ) is True
			), F"""Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})"""



def    snake_case_				(    _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE=True   ):
	model.train()
	__lowercase		       =			model(_SCREAMING_SNAKE_CASE   )
	__lowercase		       =			F.mse_loss(_SCREAMING_SNAKE_CASE  ,	target.to(output.device   )   )
	if not do_backward:
		loss /= accelerator.gradient_accumulation_steps
		loss.backward()
	else:
		accelerator.backward(_SCREAMING_SNAKE_CASE   )



def    snake_case_				(    _SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE=False   ):
	set_seed(4_2   )
	__lowercase		       =			RegressionModel()
	__lowercase		       =			deepcopy(_SCREAMING_SNAKE_CASE   )
	__lowercase		       =			RegressionDataset(length=8_0   )
	__lowercase		       =			DataLoader(_SCREAMING_SNAKE_CASE  ,	batch_size=1_6   )
	model.to(accelerator.device   )
	if sched:
		__lowercase		       =			AdamW(params=model.parameters()  ,	lr=1E-3   )
		__lowercase		       =			AdamW(params=ddp_model.parameters()  ,	lr=1E-3   )
		__lowercase		       =			LambdaLR(_SCREAMING_SNAKE_CASE  ,	lr_lambda=lambda _SCREAMING_SNAKE_CASE   : epoch**0.6_5   )
		__lowercase		       =			LambdaLR(_SCREAMING_SNAKE_CASE  ,	lr_lambda=lambda _SCREAMING_SNAKE_CASE   : epoch**0.6_5   )
	# Make a copy of `model`
	if sched:
		__lowercase       ,    __lowercase       ,    __lowercase       ,    __lowercase		       =			accelerator.prepare(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
	else:
		__lowercase       ,    __lowercase		       =			accelerator.prepare(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
	if sched:
		return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
	return model, ddp_model, dataloader



def    snake_case_				(    _SCREAMING_SNAKE_CASE   ):
	# Test when on a single CPU or GPU that the context manager does nothing
	__lowercase       ,    __lowercase       ,    __lowercase		       =			get_training_setup(_SCREAMING_SNAKE_CASE   )
	# Use a single batch
	__lowercase       ,    __lowercase		       =			next(iter(_SCREAMING_SNAKE_CASE   )   ).values()
	for iteration in range(3   ):
		# Gather the distributed inputs and targs for the base model
		__lowercase       ,    __lowercase		       =			accelerator.gather((ddp_input, ddp_target)   )
		__lowercase       ,    __lowercase		       =			input.to(accelerator.device   ), target.to(accelerator.device   )
		# Perform our initial ground truth step in non "DDP"
		step_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
		# Do "gradient accumulation" (noop)
		if iteration % 2 == 0:
			# Accumulate grads locally
			with accelerator.no_sync(_SCREAMING_SNAKE_CASE   ):
				step_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
		else:
			# Sync grads
			step_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )

		# Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
		check_model_parameters(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
		for param, ddp_param in zip(model.parameters()  ,	ddp_model.parameters()   ):
			if not param.requires_grad:
				continue
			assert torch.allclose(
			    param.grad  ,	ddp_param.grad   ), F"""Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""

		# Shuffle ddp_input on each iteration
		torch.manual_seed(1_3_3_7 + iteration   )
		__lowercase		       =			ddp_input[torch.randperm(len(_SCREAMING_SNAKE_CASE   )   )]



def    snake_case_				(    _SCREAMING_SNAKE_CASE   ):
	# Test on distributed setup that context manager behaves properly
	__lowercase       ,    __lowercase       ,    __lowercase		       =			get_training_setup(_SCREAMING_SNAKE_CASE   )
	# Use a single batch
	__lowercase       ,    __lowercase		       =			next(iter(_SCREAMING_SNAKE_CASE   )   ).values()
	for iteration in range(3   ):
		# Gather the distributed inputs and targs for the base model
		__lowercase       ,    __lowercase		       =			accelerator.gather((ddp_input, ddp_target)   )
		__lowercase       ,    __lowercase		       =			input.to(accelerator.device   ), target.to(accelerator.device   )
		# Perform our initial ground truth step in non "DDP"
		step_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
		# Do "gradient accumulation" (noop)
		if iteration % 2 == 0:
			# Accumulate grads locally
			with accelerator.no_sync(_SCREAMING_SNAKE_CASE   ):
				step_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
		else:
			# Sync grads
			step_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )

		# DDP model and model should only be in sync when not (iteration % 2 == 0)
		for param, ddp_param in zip(model.parameters()  ,	ddp_model.parameters()   ):
			if not param.requires_grad:
				continue
			if iteration % 2 == 0:
				# Grads should not be in sync
				assert (
				    torch.allclose(param.grad  ,	ddp_param.grad   ) is False
				), F"""Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"""
			else:
				# Grads should be in sync
				assert (
				    torch.allclose(param.grad  ,	ddp_param.grad   ) is True
				), F"""Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""

        # Shuffle ddp_input on each iteration
		torch.manual_seed(1_3_3_7 + iteration   )
		__lowercase		       =			ddp_input[torch.randperm(len(_SCREAMING_SNAKE_CASE   )   )]



def    snake_case_				(    _SCREAMING_SNAKE_CASE=False  ,	_SCREAMING_SNAKE_CASE=False   ):
	__lowercase		       =			Accelerator(
	    split_batches=_SCREAMING_SNAKE_CASE  ,	dispatch_batches=_SCREAMING_SNAKE_CASE  ,	gradient_accumulation_steps=2   )
	# Test that context manager behaves properly
	__lowercase       ,    __lowercase       ,    __lowercase		       =			get_training_setup(_SCREAMING_SNAKE_CASE   )
	for iteration, batch in enumerate(_SCREAMING_SNAKE_CASE   ):
		__lowercase       ,    __lowercase		       =			batch.values()
		# Gather the distributed inputs and targs for the base model
		__lowercase       ,    __lowercase		       =			accelerator.gather((ddp_input, ddp_target)   )
		__lowercase       ,    __lowercase		       =			input.to(accelerator.device   ), target.to(accelerator.device   )
		# Perform our initial ground truth step in non "DDP"
		step_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
		# Do "gradient accumulation" (noop)
		with accelerator.accumulate(_SCREAMING_SNAKE_CASE   ):
			step_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )

		# DDP model and model should only be in sync when not (iteration % 2 == 0)
		for param, ddp_param in zip(model.parameters()  ,	ddp_model.parameters()   ):
			if not param.requires_grad:
				continue
			if ((iteration + 1) % 2 == 0) or (iteration == len(_SCREAMING_SNAKE_CASE   ) - 1):
				# Grads should be in sync
				assert (
				    torch.allclose(param.grad  ,	ddp_param.grad   ) is True
				), F"""Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})"""
			else:
				# Grads should not be in sync
				assert (
				    torch.allclose(param.grad  ,	ddp_param.grad   ) is False
				), F"""Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})"""

        # Shuffle ddp_input on each iteration
		torch.manual_seed(1_3_3_7 + iteration   )
		__lowercase		       =			ddp_input[torch.randperm(len(_SCREAMING_SNAKE_CASE   )   )]
	GradientState._reset_state()



def    snake_case_				(    _SCREAMING_SNAKE_CASE=False  ,	_SCREAMING_SNAKE_CASE=False   ):
	__lowercase		       =			Accelerator(
	    split_batches=_SCREAMING_SNAKE_CASE  ,	dispatch_batches=_SCREAMING_SNAKE_CASE  ,	gradient_accumulation_steps=2   )
	# Test that context manager behaves properly
	__lowercase       ,    __lowercase       ,    __lowercase       ,    __lowercase       ,    __lowercase       ,    __lowercase       ,    __lowercase		       =			get_training_setup(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
	for iteration, batch in enumerate(_SCREAMING_SNAKE_CASE   ):
		__lowercase       ,    __lowercase		       =			batch.values()
		# Gather the distributed inputs and targs for the base model
		__lowercase       ,    __lowercase		       =			accelerator.gather((ddp_input, ddp_target)   )
		__lowercase       ,    __lowercase		       =			input.to(accelerator.device   ), target.to(accelerator.device   )
		# Perform our initial ground truth step in non "DDP"
		model.train()
		ddp_model.train()
		step_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
		opt.step()

		if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(_SCREAMING_SNAKE_CASE   )):
			if split_batches:
				sched.step()
			else:
				for _ in range(accelerator.num_processes   ):
					sched.step()
		opt.zero_grad()
		# Perform gradient accumulation under wrapper
		with accelerator.accumulate(_SCREAMING_SNAKE_CASE   ):
			step_model(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
			ddp_opt.step()
			ddp_sched.step()
			ddp_opt.zero_grad()

		# Learning rates should be the same
		assert (
		    opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
		), F"""Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]['lr']}\nDDP opt: {ddp_opt.param_groups[0]['lr']}\n"""
		__lowercase		       =			(((iteration + 1) % 2) == 0) or ((iteration + 1) == len(_SCREAMING_SNAKE_CASE   ))
		if accelerator.num_processes > 1:
			check_model_parameters(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
		# Shuffle ddp_input on each iteration
		torch.manual_seed(1_3_3_7 + iteration   )
	GradientState._reset_state()



def    snake_case_				(    ):
	__lowercase		       =			Accelerator()

	__lowercase		       =			RegressionDataset(length=8_0   )
	__lowercase		       =			DataLoader(_SCREAMING_SNAKE_CASE  ,	batch_size=1_6   )
	__lowercase		       =			RegressionDataset(length=9_6   )
	__lowercase		       =			DataLoader(_SCREAMING_SNAKE_CASE  ,	batch_size=1_6   )
	__lowercase       ,    __lowercase		       =			accelerator.prepare(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )
	assert accelerator.gradient_state.active_dataloader is None
	for iteration, _ in enumerate(_SCREAMING_SNAKE_CASE   ):
		assert id(accelerator.gradient_state.active_dataloader   ) == id(_SCREAMING_SNAKE_CASE   )
		if iteration < len(_SCREAMING_SNAKE_CASE   ) - 1:
			assert not accelerator.gradient_state.end_of_dataloader
			if iteration == 1:
				for batch_num, _ in enumerate(_SCREAMING_SNAKE_CASE   ):
					assert id(accelerator.gradient_state.active_dataloader   ) == id(_SCREAMING_SNAKE_CASE   )
					if batch_num < len(_SCREAMING_SNAKE_CASE   ) - 1:
						assert not accelerator.gradient_state.end_of_dataloader
					else:
						assert accelerator.gradient_state.end_of_dataloader
		else:
			assert accelerator.gradient_state.end_of_dataloader
	assert accelerator.gradient_state.active_dataloader is None



def    snake_case_				(    ):
	__lowercase		       =			Accelerator()
	__lowercase		       =			accelerator.state
	if state.local_process_index == 0:
		print("**Test `accumulate` gradient accumulation with dataloader break**"   )
	test_dataloader_break()
	if state.distributed_type == DistributedType.NO:
		if state.local_process_index == 0:
			print("**Test NOOP `no_sync` context manager**"   )
		test_noop_sync(_SCREAMING_SNAKE_CASE   )
	if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
		if state.local_process_index == 0:
			print("**Test Distributed `no_sync` context manager**"   )
		test_distributed_sync(_SCREAMING_SNAKE_CASE   )
	if state.distributed_type == DistributedType.MULTI_GPU:
		for split_batch in [True, False]:
			for dispatch_batches in [True, False]:
				if state.local_process_index == 0:
					print(
					    "**Test `accumulate` gradient accumulation, "  ,	F"""`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"""  ,	)
				test_gradient_accumulation(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )

    # Currently will break on torch 2.0 +, need to investigate why
	if is_torch_version("<"  ,	"2.0"   ) or state.distributed_type == DistributedType.NO:
		if state.local_process_index == 0:
			print(
			    "**Test `accumulate` gradient accumulation with optimizer and scheduler, "  ,	"`split_batches=False`, `dispatch_batches=False`**"  ,	)
		test_gradient_accumulation_with_opt_and_scheduler()
		if state.distributed_type == DistributedType.MULTI_GPU:
			for split_batch in [True, False]:
				for dispatch_batches in [True, False]:
					if not split_batch and not dispatch_batches:
						continue
					if state.local_process_index == 0:
						print(
						    "**Test `accumulate` gradient accumulation with optimizer and scheduler, "  ,	F"""`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**"""  ,	)
					test_gradient_accumulation_with_opt_and_scheduler(_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE   )



def    snake_case_				(    _SCREAMING_SNAKE_CASE   ):
	# For xla_spawn (TPUs)
	main()


if __name__ == "__main__":
	main()



