
"""simple docstring"""





import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment
def 				A_  ():

					'''simple docstring'''





					A_								=    torch.nn.Linear(2					, 4					)
					A_								=    torch.optim.AdamW(model.parameters()					, lr=1.0					)
					A_								=    torch.optim.lr_scheduler.OneCycleLR(__a					, max_lr=0.01					, steps_per_epoch=2					, epochs=1					)
					A_								=    DataLoader(TensorDataset(torch.tensor([1, 2, 3]					)					)					)
					A_								=    DataLoader(TensorDataset(torch.tensor([4, 5, 6]					)					)					)

					return model, optimizer, scheduler, train_dl, valid_dl
def 				A_  (__a					):

					'''simple docstring'''





					return (model.weight.abs().sum() + model.bias.abs().sum()).item()
def 				A_  (__a					):

					'''simple docstring'''





					A_								=    torch.nn.Linear(*tuple(model.weight.T.shape					)					).state_dict()
					model.load_state_dict(__a					)




class __lowerCAmelCase					(       _lowercase   ):




			"""simple docstring"""







			@require_cuda
			def   lowerCamelCase__  (					self    :       Any						)   ->			str:







								"""simple docstring"""
								A_								=    Accelerator()
								assert PartialState._shared_state["_cpu"] is False
								assert PartialState._shared_state["device"].type == "cuda"
								with self.assertRaises(_snake_case						):
													A_								=    Accelerator(cpu=_snake_case						)




			def   lowerCamelCase__  (					self    :       int						)   ->			List[str]:







								"""simple docstring"""
								A_								=    Accelerator()
								A_								=    GradientState()
								assert state.num_steps == 1
								A_								=    4
								assert state.num_steps == 4

								assert state.sync_gradients is True
								A_								=    False
								assert state.sync_gradients is False
								GradientState._reset_state()




			def   lowerCamelCase__  (					self    :       int						)   ->			Optional[int]:







								"""simple docstring"""
								A_								=    Accelerator()
								A_			,     A_			,     A_			,     A_			,     A_								=    create_components()

								(
								    (
								    A_
								)			,     (
								    A_
								)			,     (
								    A_
								)			,     (
								    A_
								)			,     (
								    A_
								)			,     
								)								=    accelerator.prepare(_snake_case							, _snake_case							, _snake_case							, _snake_case							, _snake_case						)

								self.assertTrue(prepared_model in accelerator._models						)
								self.assertTrue(prepared_optimizer in accelerator._optimizers						)
								self.assertTrue(prepared_scheduler in accelerator._schedulers						)
								self.assertTrue(prepared_train_dl in accelerator._dataloaders						)
								self.assertTrue(prepared_valid_dl in accelerator._dataloaders						)




			def   lowerCamelCase__  (					self    :       int						)   ->			Optional[Any]:







								"""simple docstring"""
								A_								=    Accelerator()
								A_			,     A_			,     A_			,     A_			,     A_								=    create_components()
								accelerator.prepare(_snake_case							, _snake_case							, _snake_case							, _snake_case							, _snake_case						)
								accelerator.free_memory()

								self.assertTrue(len(accelerator._models						) == 0						)
								self.assertTrue(len(accelerator._optimizers						) == 0						)
								self.assertTrue(len(accelerator._schedulers						) == 0						)
								self.assertTrue(len(accelerator._dataloaders						) == 0						)




			def   lowerCamelCase__  (					self    :       Dict						)   ->			List[Any]:







								"""simple docstring"""
								PartialState._reset_state()

								# Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
								def noop(*_snake_case    :       Tuple							, **_snake_case    :       Dict						):
													pass

								with patch("torch.cuda.set_device"							, _snake_case						), patch_environment(ACCELERATE_TORCH_DEVICE="cuda:64"						):
													A_								=    Accelerator()
													self.assertEqual(str(accelerator.state.device						)							, "cuda:64"						)




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







								"""simple docstring"""
								A_								=    Accelerator()
								A_			,     A_			,     A_			,     A_			,     A_								=    create_components()
								accelerator.prepare(_snake_case							, _snake_case							, _snake_case							, _snake_case							, _snake_case						)

								A_								=    get_signature(_snake_case						)

								with tempfile.TemporaryDirectory() as tmpdirname:
													accelerator.save_state(_snake_case						)

													# make sure random weights don't match
													load_random_weights(_snake_case						)
													self.assertTrue(abs(model_signature - get_signature(_snake_case						)						) > 1e-3						)

													# make sure loaded weights match
													accelerator.load_state(_snake_case						)
													self.assertTrue(abs(model_signature - get_signature(_snake_case						)						) < 1e-3						)




			def   lowerCamelCase__  (					self    :       Optional[Any]						)   ->			str:







								"""simple docstring"""
								A_								=    Accelerator()
								A_			,     A_			,     A_			,     A_			,     A_								=    create_components()
								accelerator.prepare(_snake_case							, _snake_case							, _snake_case							, _snake_case							, _snake_case						)

								A_								=    get_signature(_snake_case						)

								# saving hook
								def save_config(_snake_case    :       Union[str, Any]							, _snake_case    :       List[str]							, _snake_case    :       List[Any]						):
													A_								=    {"class_name": models[0].__class__.__name__}

													with open(os.path.join(_snake_case							, "data.json"						)							, "w"						) as f:
																		json.dump(_snake_case							, _snake_case						)

        # loading hook
								def load_config(_snake_case    :       Any							, _snake_case    :       List[str]						):
													with open(os.path.join(_snake_case							, "data.json"						)							, "r"						) as f:
																		A_								=    json.load(_snake_case						)

													A_								=    config["class_name"]

								A_								=    accelerator.register_save_state_pre_hook(_snake_case						)
								A_								=    accelerator.register_load_state_pre_hook(_snake_case						)

								with tempfile.TemporaryDirectory() as tmpdirname:
													accelerator.save_state(_snake_case						)

													# make sure random weights don't match with hooks
													load_random_weights(_snake_case						)
													self.assertTrue(abs(model_signature - get_signature(_snake_case						)						) > 1e-3						)

													# random class name to verify correct one is loaded
													A_								=    "random"

													# make sure loaded weights match with hooks
													accelerator.load_state(_snake_case						)
													self.assertTrue(abs(model_signature - get_signature(_snake_case						)						) < 1e-3						)

													# mode.class_name is loaded from config
													self.assertTrue(model.class_name == model.__class__.__name__						)

								# remove hooks
								save_hook.remove()
								load_hook.remove()

								with tempfile.TemporaryDirectory() as tmpdirname:
													accelerator.save_state(_snake_case						)

													# make sure random weights don't match with hooks removed
													load_random_weights(_snake_case						)
													self.assertTrue(abs(model_signature - get_signature(_snake_case						)						) > 1e-3						)

													# random class name to verify correct one is loaded
													A_								=    "random"

													# make sure loaded weights match with hooks removed
													accelerator.load_state(_snake_case						)
													self.assertTrue(abs(model_signature - get_signature(_snake_case						)						) < 1e-3						)

													# mode.class_name is NOT loaded from config
													self.assertTrue(model.class_name != model.__class__.__name__						)




			def   lowerCamelCase__  (					self    :       int						)   ->			Optional[Any]:







								"""simple docstring"""
								A_								=    Accelerator()
								A_			,     A_			,     A_			,     A_			,     A_								=    create_components()
								A_								=    None

								# This should work
								A_			,     A_			,     A_			,     A_			,     A_			,     A_								=    accelerator.prepare(
								    _snake_case							, _snake_case							, _snake_case							, _snake_case							, _snake_case							, _snake_case						)
								self.assertTrue(dummy_obj is None						)




			def   lowerCamelCase__  (					self    :       List[str]						)   ->			List[str]:







								"""simple docstring"""
								A_								=    Accelerator()
								A_			,     A_			,     A_			,     A_			,     A_								=    create_components()
								A_								=    [1, 2, 3]

								# This should work
								A_			,     A_			,     A_			,     A_			,     A_			,     A_								=    accelerator.prepare(
								    _snake_case							, _snake_case							, _snake_case							, _snake_case							, _snake_case							, _snake_case						)
								self.assertEqual(
								    getattr(_snake_case							, "_is_accelerate_prepared"							, _snake_case						)							, _snake_case							, "Dummy object should have `_is_accelerate_prepared` set to `True`"							, )
								self.assertEqual(
								    getattr(_snake_case							, "_is_accelerate_prepared"							, _snake_case						)							, _snake_case							, "Model is missing `_is_accelerator_prepared` or is set to `False`"							, )
								self.assertEqual(
								    getattr(_snake_case							, "_is_accelerate_prepared"							, _snake_case						)							, _snake_case							, "Optimizer is missing `_is_accelerator_prepared` or is set to `False`"							, )
								self.assertEqual(
								    getattr(_snake_case							, "_is_accelerate_prepared"							, _snake_case						)							, _snake_case							, "Scheduler is missing `_is_accelerator_prepared` or is set to `False`"							, )
								self.assertEqual(
								    getattr(_snake_case							, "_is_accelerate_prepared"							, _snake_case						)							, _snake_case							, "Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`"							, )
								self.assertEqual(
								    getattr(_snake_case							, "_is_accelerate_prepared"							, _snake_case						)							, _snake_case							, "Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`"							, )




			@slow
			@require_bnb
			def   lowerCamelCase__  (					self    :       Any						)   ->			str:







								"""simple docstring"""
								from transformers import AutoModelForCausalLM

								A_								=    AutoModelForCausalLM.from_pretrained(
								    "EleutherAI/gpt-neo-125m"							, load_in_abit=_snake_case							, device_map={"": 0}							, )
								A_								=    Accelerator()

								# This should work
								A_								=    accelerator.prepare(_snake_case						)




			@slow
			@require_bnb
			def   lowerCamelCase__  (					self    :       Tuple						)   ->			Optional[Any]:







								"""simple docstring"""
								from transformers import AutoModelForCausalLM

								A_								=    Accelerator()

								with init_empty_weights():
													A_								=    AutoModelForCausalLM.from_pretrained(
													    "EleutherAI/gpt-neo-125m"							, )
													model.tie_weights()
													A_								=    infer_auto_device_map(_snake_case						)
													A_								=    "cpu"

								A_								=    AutoModelForCausalLM.from_pretrained(
								    "EleutherAI/gpt-neo-125m"							, device_map=_snake_case							, load_in_abit=_snake_case							, llm_inta_enable_fpaa_cpu_offload=_snake_case						)

								# This should not work and get value error
								with self.assertRaises(_snake_case						):
													A_								=    accelerator.prepare(_snake_case						)




			@slow
			@require_bnb
			@require_multi_gpu
			def   lowerCamelCase__  (					self    :       Union[str, Any]						)   ->			Optional[int]:







								"""simple docstring"""
								from transformers import AutoModelForCausalLM

								A_								=    {"distributed_type": DistributedType.MULTI_GPU}

								with init_empty_weights():
													A_								=    AutoModelForCausalLM.from_pretrained(
													    "EleutherAI/gpt-neo-125m"							, )
													model.tie_weights()
													A_								=    infer_auto_device_map(_snake_case						)
													A_								=    1

								A_								=    AutoModelForCausalLM.from_pretrained(
								    "EleutherAI/gpt-neo-125m"							, load_in_abit=_snake_case							, device_map=_snake_case							, )
								A_								=    Accelerator()

								# This should not work and get value error
								with self.assertRaises(_snake_case						):
													A_								=    accelerator.prepare(_snake_case						)

								PartialState._reset_state()




			@slow
			@require_bnb
			@require_multi_gpu
			def   lowerCamelCase__  (					self    :       List[str]						)   ->			Optional[Any]:







								"""simple docstring"""
								from transformers import AutoModelForCausalLM

								with init_empty_weights():
													A_								=    AutoModelForCausalLM.from_pretrained(
													    "EleutherAI/gpt-neo-125m"							, )
													A_								=    infer_auto_device_map(_snake_case						)
													A_								=    1

								A_								=    AutoModelForCausalLM.from_pretrained(
								    "EleutherAI/gpt-neo-125m"							, load_in_abit=_snake_case							, device_map=_snake_case							, )
								A_								=    Accelerator()

								# This should work
								A_								=    accelerator.prepare(_snake_case						)






			@require_cuda
			def   lowerCamelCase__  (					self    :       Any						)   ->			str:







								"""simple docstring"""
								A_								=    torch.nn.Linear(10							, 10						)
								A_								=    torch.optim.SGD(model.parameters()							, lr=0.0_1						)
								A_								=    Accelerator(cpu=_snake_case						)
								A_								=    accelerator.prepare(_snake_case						)






"""simple docstring"""





import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
		import os

		import jax.numpy as jnp
		from jax import jit

		from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
		from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

		UpperCamelCase_       :    List[Any]						      =			'''0.12'''  # assumed parallelism: 8


if is_torch_available():
		import torch
def 				A_  (__a					, __a					, __a=None					):

					'''simple docstring'''





					if rng is None:
										A_								=    random.Random()

					A_								=    1
					for dim in shape:
										total_dims *= dim

					A_								=    []
					for _ in range(__a					):
										values.append(rng.randint(0					, vocab_size - 1					)					)

					A_								=    np.array(__a					, dtype=jnp.intaa					).reshape(__a					)

					return output
def 				A_  (__a					, __a=None					):

					'''simple docstring'''





					A_								=    ids_tensor(__a					, vocab_size=2					, rng=__a					)
					# make sure that at least one token is attended to for each batch
					A_								=    1
					return attn_mask



@require_flax
class __lowerCAmelCase					:




			"""simple docstring"""







			snake_case						   =					None
			snake_case						   =					()




			def   lowerCamelCase__  (					self    :       Any						)   ->			Tuple:







								"""simple docstring"""
								A_			,     A_								=    self.model_tester.prepare_config_and_inputs_for_common()

								# cut to half length & take max batch_size 3
								A_								=    2
								A_								=    inputs["input_ids"].shape[-1] // 2
								A_								=    inputs["input_ids"][:max_batch_size, :sequence_length]

								A_								=    jnp.ones_like(_snake_case						)
								A_								=    attention_mask[:max_batch_size, :sequence_length]

								# generate max 5 tokens
								A_								=    input_ids.shape[-1] + 5
								if config.eos_token_id is not None and config.pad_token_id is None:
													# hack to allow generate for models such as GPT2 as is done in `generate()`
													A_								=    config.eos_token_id
								return config, input_ids, attention_mask, max_length




			@is_pt_flax_cross_test
			def   lowerCamelCase__  (					self    :       List[Any]						)   ->			List[Any]:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length
								A_								=    0

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model_class.__name__[4:]  # Skip the "Flax" at the beginning
													A_								=    getattr(_snake_case							, _snake_case						)
													A_								=    pt_model_class(_snake_case						).eval()
													A_								=    load_flax_weights_in_pytorch_model(_snake_case							, flax_model.params						)

													A_								=    flax_model.generate(_snake_case						).sequences
													A_								=    pt_model.generate(torch.tensor(_snake_case							, dtype=torch.long						)						)

													if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
																		A_								=    flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

													self.assertListEqual(pt_generation_outputs.numpy().tolist()							, flax_generation_outputs.tolist()						)




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







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			List[str]:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    True
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			Dict:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length
								A_								=    2

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       str						)   ->			Union[str, Any]:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    False
								A_								=    max_length
								A_								=    2
								A_								=    2

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[0]							, input_ids.shape[0] * config.num_return_sequences						)




			def   lowerCamelCase__  (					self    :       Union[str, Any]						)   ->			Dict:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    True
								A_								=    max_length
								A_								=    0.8
								A_								=    10
								A_								=    0.3
								A_								=    1
								A_								=    8
								A_								=    9

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       List[str]						)   ->			Any:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    max_length
								A_								=    1
								A_								=    8
								A_								=    9

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Optional[Any]						)   ->			Dict:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()
								A_								=    max_length
								A_								=    2
								A_								=    1
								A_								=    8
								A_								=    9

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       Tuple						)   ->			Any:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()

								# pad attention mask on the left
								A_								=    attention_mask.at[(0, 0)].set(0						)

								A_								=    False
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case							, attention_mask=_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case							, attention_mask=_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)




			def   lowerCamelCase__  (					self    :       List[str]						)   ->			int:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()

								# pad attention mask on the left
								A_								=    attention_mask.at[(0, 0)].set(0						)

								A_								=    True
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case							, attention_mask=_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case							, attention_mask=_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)






			def   lowerCamelCase__  (					self    :       Dict						)   ->			int:







								"""simple docstring"""
								A_			,     A_			,     A_			,     A_								=    self._get_input_ids_and_config()

								# pad attention mask on the left
								A_								=    attention_mask.at[(0, 0)].set(0						)

								A_								=    2
								A_								=    max_length

								for model_class in self.all_generative_model_classes:
													A_								=    model_class(_snake_case						)

													A_								=    model.generate(_snake_case							, attention_mask=_snake_case						).sequences
													self.assertEqual(generation_outputs.shape[-1]							, _snake_case						)

													A_								=    jit(model.generate						)
													A_								=    jit_generate(_snake_case							, attention_mask=_snake_case						).sequences

													self.assertListEqual(generation_outputs.tolist()							, jit_generation_outputs.tolist()						)





@require_flax
class __lowerCAmelCase					(       unittest.TestCase   ):




			"""simple docstring"""







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







								"""simple docstring"""
								A_								=    AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"						)
								A_								=    FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"						)

								A_								=    "Hello world"
								A_								=    tokenizer(_snake_case							, return_tensors="np"						).input_ids

								# typos are quickly detected (the correct argument is `do_sample`)
								with self.assertRaisesRegex(_snake_case							, "do_samples"						):
													model.generate(_snake_case							, do_samples=_snake_case						)

								# arbitrary arguments that will not be used anywhere are also not accepted
								with self.assertRaisesRegex(_snake_case							, "foo"						):
													A_								=    {"foo": "bar"}
													model.generate(_snake_case							, **_snake_case						)




