






"""simple docstring"""


import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
		import jax

		from transformers.models.big_bird.modeling_flax_big_bird import (
		    FlaxBigBirdForCausalLM,
		    FlaxBigBirdForMaskedLM,
		    FlaxBigBirdForMultipleChoice,
		    FlaxBigBirdForPreTraining,
		    FlaxBigBirdForQuestionAnswering,
		    FlaxBigBirdForSequenceClassification,
		    FlaxBigBirdForTokenClassification,
		    FlaxBigBirdModel,
		)




class 						__UpperCamelCase     (     unittest.TestCase						):
					def __init__(							self							:	Tuple     , UpperCAmelCase							:	List[Any]     , UpperCAmelCase							:	Any=2     , UpperCAmelCase							:	Any=56     , UpperCAmelCase							:	Dict=True     , UpperCAmelCase							:	str=True     , UpperCAmelCase							:	Optional[int]=True     , UpperCAmelCase							:	Union[str, Any]=True     , UpperCAmelCase							:	Tuple=99     , UpperCAmelCase							:	int=32     , UpperCAmelCase							:	Any=2     , UpperCAmelCase							:	Dict=2     , UpperCAmelCase							:	Union[str, Any]=7     , UpperCAmelCase							:	Optional[Any]="gelu_new"     , UpperCAmelCase							:	List[Any]=0.1     , UpperCAmelCase							:	List[Any]=0.1     , UpperCAmelCase							:	List[str]=512     , UpperCAmelCase							:	List[Any]=16     , UpperCAmelCase							:	List[str]=2     , UpperCAmelCase							:	List[Any]=0.0_2     , UpperCAmelCase							:	Optional[int]=4     , UpperCAmelCase							:	Union[str, Any]="block_sparse"     , UpperCAmelCase							:	Any=True     , UpperCAmelCase							:	Union[str, Any]=False     , UpperCAmelCase							:	List[Any]=2     , UpperCAmelCase							:	Optional[Any]=3     , )	->       int:
											lowerCAmelCase   :Optional[int]						=			parent
											lowerCAmelCase   :Dict						=			batch_size
											lowerCAmelCase   :Any						=			seq_length
											lowerCAmelCase   :Optional[int]						=			is_training
											lowerCAmelCase   :List[str]						=			use_attention_mask
											lowerCAmelCase   :Dict						=			use_token_type_ids
											lowerCAmelCase   :int						=			use_labels
											lowerCAmelCase   :Optional[int]						=			vocab_size
											lowerCAmelCase   :Any						=			hidden_size
											lowerCAmelCase   :Dict						=			num_hidden_layers
											lowerCAmelCase   :List[Any]						=			num_attention_heads
											lowerCAmelCase   :Optional[Any]						=			intermediate_size
											lowerCAmelCase   :int						=			hidden_act
											lowerCAmelCase   :Optional[Any]						=			hidden_dropout_prob
											lowerCAmelCase   :Optional[Any]						=			attention_probs_dropout_prob
											lowerCAmelCase   :int						=			max_position_embeddings
											lowerCAmelCase   :Optional[Any]						=			type_vocab_size
											lowerCAmelCase   :Any						=			type_sequence_label_size
											lowerCAmelCase   :Tuple						=			initializer_range
											lowerCAmelCase   :Optional[int]						=			num_choices

											lowerCAmelCase   :Optional[int]						=			rescale_embeddings
											lowerCAmelCase   :Any						=			attention_type
											lowerCAmelCase   :str						=			use_bias
											lowerCAmelCase   :Any						=			block_size
											lowerCAmelCase   :str						=			num_random_blocks
					def        UpperCAmelCase__ (							self							:	Dict   )	->       Optional[int]:
											lowerCAmelCase   :Optional[int]						=			ids_tensor([self.batch_size, self.seq_length]     , self.vocab_size   )

											lowerCAmelCase   :Tuple						=			None
											if self.use_attention_mask:
																	lowerCAmelCase   :List[Any]						=			random_attention_mask([self.batch_size, self.seq_length]   )

											lowerCAmelCase   :Optional[int]						=			None
											if self.use_token_type_ids:
																	lowerCAmelCase   :str						=			ids_tensor([self.batch_size, self.seq_length]     , self.type_vocab_size   )

											lowerCAmelCase   :Dict						=			BigBirdConfig(
											    vocab_size=self.vocab_size     , hidden_size=self.hidden_size     , num_hidden_layers=self.num_hidden_layers     , num_attention_heads=self.num_attention_heads     , intermediate_size=self.intermediate_size     , hidden_act=self.hidden_act     , hidden_dropout_prob=self.hidden_dropout_prob     , attention_probs_dropout_prob=self.attention_probs_dropout_prob     , max_position_embeddings=self.max_position_embeddings     , type_vocab_size=self.type_vocab_size     , is_decoder=UpperCAmelCase     , initializer_range=self.initializer_range     , attention_type=self.attention_type     , block_size=self.block_size     , num_random_blocks=self.num_random_blocks     , use_bias=self.use_bias     , rescale_embeddings=self.rescale_embeddings     , )

											return config, input_ids, token_type_ids, attention_mask
					def        UpperCAmelCase__ (							self							:	Dict   )	->       Optional[int]:
											lowerCAmelCase   :Optional[int]						=			self.prepare_config_and_inputs()
											lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase   :List[str]						=			config_and_inputs
											lowerCAmelCase   :Any						=			{
											    'input_ids': input_ids,
											    'token_type_ids': token_type_ids,
											    'attention_mask': attention_mask,
											}
											return config, inputs_dict




@require_flax
class 						__UpperCamelCase     (     UpperCamelCase   ,     unittest.TestCase						):
					lowercase_    :				List[Any]      =  (
					    (
					        FlaxBigBirdForCausalLM,
					        FlaxBigBirdModel,
					        FlaxBigBirdForPreTraining,
					        FlaxBigBirdForMaskedLM,
					        FlaxBigBirdForMultipleChoice,
					        FlaxBigBirdForQuestionAnswering,
					        FlaxBigBirdForSequenceClassification,
					        FlaxBigBirdForTokenClassification,
					    )
					    if is_flax_available()
					    else ()
					)

					lowercase_    :				Optional[int]      =  False
					lowercase_    :				Tuple      =  False
					def        UpperCAmelCase__ (							self							:	List[Any]   )	->       Optional[int]:
											lowerCAmelCase   :List[Any]						=			FlaxBigBirdModelTester(self   )
					@slow
					# copied from `test_modeling_flax_common` because it takes much longer than other models
					def        UpperCAmelCase__ (							self							:	Any   )	->       Tuple:
											super().test_from_pretrained_save_pretrained()
					@slow
					# copied from `test_modeling_flax_common` because it takes much longer than other models
					def        UpperCAmelCase__ (							self							:	str   )	->       Tuple:
											super().test_from_pretrained_with_no_automatic_init()
					@slow
					# copied from `test_modeling_flax_common` because it takes much longer than other models
					def        UpperCAmelCase__ (							self							:	Dict   )	->       Optional[int]:
											super().test_no_automatic_init()
					@slow
					# copied from `test_modeling_flax_common` because it takes much longer than other models
					def        UpperCAmelCase__ (							self							:	Union[str, Any]   )	->       str:
											super().test_hidden_states_output()
					@slow
					def        UpperCAmelCase__ (							self							:	Tuple   )	->       Dict:
											for model_class_name in self.all_model_classes:
																	lowerCAmelCase   :Dict						=			model_class_name.from_pretrained('google/bigbird-roberta-base'   )
																	self.assertIsNotNone(UpperCAmelCase   )
					def        UpperCAmelCase__ (							self							:	Optional[int]   )	->       List[str]:
											if self.test_attn_probs:
																	super().test_attention_outputs()
					@slow
					# copied from `test_modeling_flax_common` because it takes much longer than other models
					def        UpperCAmelCase__ (							self							:	Optional[int]   )	->       Any:
											lowerCAmelCase	, lowerCAmelCase   :List[str]						=			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	with self.subTest(model_class.__name__   ):
																							lowerCAmelCase   :int						=			self._prepare_for_class(UpperCAmelCase     , UpperCAmelCase   )
																							lowerCAmelCase   :Tuple						=			model_class(UpperCAmelCase   )

																							@jax.jit
																							def model_jitted(UpperCAmelCase							:	str     , UpperCAmelCase							:	Any=None     , **UpperCAmelCase							:	Optional[int]   ):
																													return model(input_ids=UpperCAmelCase     , attention_mask=UpperCAmelCase     , **UpperCAmelCase   )

																							with self.subTest('JIT Enabled'   ):
																													lowerCAmelCase   :Optional[int]						=			model_jitted(**UpperCAmelCase   ).to_tuple()

																							with self.subTest('JIT Disabled'   ):
																													with jax.disable_jit():
																																			lowerCAmelCase   :int						=			model_jitted(**UpperCAmelCase   ).to_tuple()

																							self.assertEqual(len(UpperCAmelCase   )     , len(UpperCAmelCase   )   )
																							for jitted_output, output in zip(UpperCAmelCase     , UpperCAmelCase   ):
																													self.assertEqual(jitted_output.shape     , output.shape   )
					def        UpperCAmelCase__ (							self							:	Optional[Any]     , UpperCAmelCase							:	Any     , UpperCAmelCase							:	List[Any]     , UpperCAmelCase							:	Any     , UpperCAmelCase							:	str=1e-5     , UpperCAmelCase							:	Tuple="outputs"     , UpperCAmelCase							:	Any=None   )	->       Optional[int]:
											# `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
											# an effort was done to return `attention_probs` (yet to be verified).
											if name.startswith('outputs.attentions'   ):
																	return
											else:
																	super().check_pt_flax_outputs(UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase     , UpperCAmelCase   )







"""simple docstring"""


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        UpperCAmelCase      (					a__     ,      a__     ,      a__     ,      a__ ):





						'''simple docstring'''






						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        UpperCAmelCase      (					a__     ,      a__     ,      a__     ,      a__     ,      a__=True ):





						'''simple docstring'''






						model.train()
						lowerCAmelCase   :Dict						=			model(a__ )
						lowerCAmelCase   :Union[str, Any]						=			F.mse_loss(a__     ,      target.to(output.device ) )
						if not do_backward:
												loss /= accelerator.gradient_accumulation_steps
												loss.backward()
						else:
												accelerator.backward(a__ )


def        UpperCAmelCase      (					a__     ,      a__=False ):





						'''simple docstring'''






						set_seed(42 )
						lowerCAmelCase   :Any						=			RegressionModel()
						lowerCAmelCase   :List[str]						=			deepcopy(a__ )
						lowerCAmelCase   :str						=			RegressionDataset(length=80 )
						lowerCAmelCase   :Optional[Any]						=			DataLoader(a__     ,      batch_size=16 )
						model.to(accelerator.device )
						if sched:
												lowerCAmelCase   :List[Any]						=			AdamW(params=model.parameters()     ,      lr=1e-3 )
												lowerCAmelCase   :Dict						=			AdamW(params=ddp_model.parameters()     ,      lr=1e-3 )
												lowerCAmelCase   :Tuple						=			LambdaLR(a__     ,      lr_lambda=lambda a__ : epoch**0.65 )
												lowerCAmelCase   :Optional[Any]						=			LambdaLR(a__     ,      lr_lambda=lambda a__ : epoch**0.65 )
						# Make a copy of `model`
						if sched:
												lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase   :int						=			accelerator.prepare(a__     ,      a__     ,      a__     ,      a__ )
						else:
												lowerCAmelCase	, lowerCAmelCase   :Optional[int]						=			accelerator.prepare(a__     ,      a__ )
						if sched:
												return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
						return model, ddp_model, dataloader


def        UpperCAmelCase      (					a__ ):





						'''simple docstring'''






						lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase   :Any						=			get_training_setup(a__ )
						# Use a single batch
						lowerCAmelCase	, lowerCAmelCase   :Any						=			next(iter(a__ ) ).values()
						for iteration in range(3 ):
												# Gather the distributed inputs and targs for the base model
												lowerCAmelCase	, lowerCAmelCase   :Optional[Any]						=			accelerator.gather((ddp_input, ddp_target) )
												lowerCAmelCase	, lowerCAmelCase   :Optional[Any]						=			input.to(accelerator.device ), target.to(accelerator.device )
												# Perform our initial ground truth step in non "DDP"
												step_model(a__     ,      a__     ,      a__     ,      a__ )
												# Do "gradient accumulation" (noop)
												if iteration % 2 == 0:
																		# Accumulate grads locally
																		with accelerator.no_sync(a__ ):
																								step_model(a__     ,      a__     ,      a__     ,      a__ )
												else:
																		# Sync grads
																		step_model(a__     ,      a__     ,      a__     ,      a__ )

												# Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
												check_model_parameters(a__     ,      a__     ,      a__     ,      a__ )
												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(13_37 + iteration )
												lowerCAmelCase   :Any						=			ddp_input[torch.randperm(len(a__ ) )]


def        UpperCAmelCase      (					a__ ):





						'''simple docstring'''






						lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase   :Dict						=			get_training_setup(a__ )
						# Use a single batch
						lowerCAmelCase	, lowerCAmelCase   :str						=			next(iter(a__ ) ).values()
						for iteration in range(3 ):
												# Gather the distributed inputs and targs for the base model
												lowerCAmelCase	, lowerCAmelCase   :List[Any]						=			accelerator.gather((ddp_input, ddp_target) )
												lowerCAmelCase	, lowerCAmelCase   :int						=			input.to(accelerator.device ), target.to(accelerator.device )
												# Perform our initial ground truth step in non "DDP"
												step_model(a__     ,      a__     ,      a__     ,      a__ )
												# Do "gradient accumulation" (noop)
												if iteration % 2 == 0:
																		# Accumulate grads locally
																		with accelerator.no_sync(a__ ):
																								step_model(a__     ,      a__     ,      a__     ,      a__ )
												else:
																		# Sync grads
																		step_model(a__     ,      a__     ,      a__     ,      a__ )

												# 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(13_37 + iteration )
												lowerCAmelCase   :str						=			ddp_input[torch.randperm(len(a__ ) )]


def        UpperCAmelCase      (					a__=False     ,      a__=False ):





						'''simple docstring'''






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

												# 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(a__ ) - 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(13_37 + iteration )
												lowerCAmelCase   :Optional[int]						=			ddp_input[torch.randperm(len(a__ ) )]
						GradientState._reset_state()


def        UpperCAmelCase      (					a__=False     ,      a__=False ):





						'''simple docstring'''






						lowerCAmelCase   :Optional[int]						=			Accelerator(
						    split_batches=a__     ,      dispatch_batches=a__     ,      gradient_accumulation_steps=2 )
						# Test that context manager behaves properly
						lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase	, lowerCAmelCase   :Union[str, Any]						=			get_training_setup(a__     ,      a__ )
						for iteration, batch in enumerate(a__ ):
												lowerCAmelCase	, lowerCAmelCase   :Optional[int]						=			batch.values()
												# Gather the distributed inputs and targs for the base model
												lowerCAmelCase	, lowerCAmelCase   :Union[str, Any]						=			accelerator.gather((ddp_input, ddp_target) )
												lowerCAmelCase	, lowerCAmelCase   :Optional[Any]						=			input.to(accelerator.device ), target.to(accelerator.device )
												# Perform our initial ground truth step in non "DDP"
												model.train()
												ddp_model.train()
												step_model(a__     ,      a__     ,      a__     ,      a__     ,      a__ )
												opt.step()

												if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(a__ )):
																		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(a__ ):
																		step_model(a__     ,      a__     ,      a__     ,      a__ )
																		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"""
												lowerCAmelCase   :int						=			(((iteration + 1) % 2) == 0) or ((iteration + 1) == len(a__ ))
												if accelerator.num_processes > 1:
																		check_model_parameters(a__     ,      a__     ,      a__     ,      a__ )
												# Shuffle ddp_input on each iteration
												torch.manual_seed(13_37 + iteration )
						GradientState._reset_state()


def        UpperCAmelCase      (					):





						'''simple docstring'''






						lowerCAmelCase   :int						=			Accelerator()

						lowerCAmelCase   :int						=			RegressionDataset(length=80 )
						lowerCAmelCase   :Optional[Any]						=			DataLoader(a__     ,      batch_size=16 )
						lowerCAmelCase   :Any						=			RegressionDataset(length=96 )
						lowerCAmelCase   :Dict						=			DataLoader(a__     ,      batch_size=16 )
						lowerCAmelCase	, lowerCAmelCase   :Tuple						=			accelerator.prepare(a__     ,      a__ )
						assert accelerator.gradient_state.active_dataloader is None
						for iteration, _ in enumerate(a__ ):
												assert id(accelerator.gradient_state.active_dataloader ) == id(a__ )
												if iteration < len(a__ ) - 1:
																		assert not accelerator.gradient_state.end_of_dataloader
																		if iteration == 1:
																								for batch_num, _ in enumerate(a__ ):
																														assert id(accelerator.gradient_state.active_dataloader ) == id(a__ )
																														if batch_num < len(a__ ) - 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        UpperCAmelCase      (					):





						'''simple docstring'''






						lowerCAmelCase   :List[str]						=			Accelerator()
						lowerCAmelCase   :Optional[int]						=			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(a__ )
						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(a__ )
						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(a__     ,      a__ )

    # 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(a__     ,      a__ )


def        UpperCAmelCase      (					a__ ):





						'''simple docstring'''






						main()


if __name__ == "__main__":
		main()