




import unittest

from transformers import (
    MODEL_FOR_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
    TextGenerationPipeline,
    logging,
    pipeline,
)
from transformers.testing_utils import (
    CaptureLogger,
    is_pipeline_test,
    require_accelerate,
    require_tf,
    require_torch,
    require_torch_gpu,
    require_torch_or_tf,
)

from .test_pipelines_common import ANY




@is_pipeline_test
@require_torch_or_tf
class    __lowerCAmelCase       (							unittest.TestCase):
       _a 					=						MODEL_FOR_CAUSAL_LM_MAPPING
       _a 					=						TF_MODEL_FOR_CAUSAL_LM_MAPPING






       @require_torch
       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[Any]				):
         lowercase    :Tuple      	=						pipeline(task="text-generation"       ,					model="sshleifer/tiny-ctrl"       ,					framework="pt"				)
         # Using `do_sample=False` to force deterministic output
         lowercase    :Dict      	=						text_generator("This is a test"       ,					do_sample=_lowerCAmelCase				)
         self.assertEqual(
             _lowerCAmelCase       ,					[
                 {
                     "generated_text": (
                         "This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope."
                         " oscope. FiliFili@@"
                     )
                 }
             ]       ,					)

         lowercase    :Tuple      	=						text_generator(["This is a test", "This is a second test"]				)
         self.assertEqual(
             _lowerCAmelCase       ,					[
                 [
                     {
                         "generated_text": (
                             "This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope."
                             " oscope. FiliFili@@"
                         )
                     }
                 ],
                 [
                     {
                         "generated_text": (
                             "This is a second test ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy"
                             " oscope. oscope. FiliFili@@"
                         )
                     }
                 ],
             ]       ,					)

         lowercase    :Optional[int]      	=						text_generator("This is a test"       ,					do_sample=_lowerCAmelCase       ,					num_return_sequences=2       ,					return_tensors=_lowerCAmelCase				)
         self.assertEqual(
             _lowerCAmelCase       ,					[
                 {"generated_token_ids": ANY(_lowerCAmelCase				)},
                 {"generated_token_ids": ANY(_lowerCAmelCase				)},
             ]       ,					)
         lowercase    :Any      	=						text_generator.model.config.eos_token_id
         lowercase    :int      	=						"<pad>"
         lowercase    :Optional[Any]      	=						text_generator(
             ["This is a test", "This is a second test"]       ,					do_sample=_lowerCAmelCase       ,					num_return_sequences=2       ,					batch_size=2       ,					return_tensors=_lowerCAmelCase       ,					)
         self.assertEqual(
             _lowerCAmelCase       ,					[
                 [
                     {"generated_token_ids": ANY(_lowerCAmelCase				)},
                     {"generated_token_ids": ANY(_lowerCAmelCase				)},
                 ],
                 [
                     {"generated_token_ids": ANY(_lowerCAmelCase				)},
                     {"generated_token_ids": ANY(_lowerCAmelCase				)},
                 ],
             ]       ,					)






       @require_tf
       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         lowercase    :Tuple      	=						pipeline(task="text-generation"       ,					model="sshleifer/tiny-ctrl"       ,					framework="tf"				)

         # Using `do_sample=False` to force deterministic output
         lowercase    :Tuple      	=						text_generator("This is a test"       ,					do_sample=_lowerCAmelCase				)
         self.assertEqual(
             _lowerCAmelCase       ,					[
                 {
                     "generated_text": (
                         "This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵"
                         " please,"
                     )
                 }
             ]       ,					)

         lowercase    :List[str]      	=						text_generator(["This is a test", "This is a second test"]       ,					do_sample=_lowerCAmelCase				)
         self.assertEqual(
             _lowerCAmelCase       ,					[
                 [
                     {
                         "generated_text": (
                             "This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵"
                             " please,"
                         )
                     }
                 ],
                 [
                     {
                         "generated_text": (
                             "This is a second test Chieftain Chieftain prefecture prefecture prefecture Cannes Cannes"
                             " Cannes 閲閲Cannes Cannes Cannes 攵 please,"
                         )
                     }
                 ],
             ]       ,					)






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]       ,					_lowerCAmelCase:					Union[str, Any]       ,					_lowerCAmelCase:					Optional[int]       ,					_lowerCAmelCase:					str				):
         lowercase    :int      	=						TextGenerationPipeline(model=_lowerCAmelCase       ,					tokenizer=_lowerCAmelCase				)
         return text_generator, ["This is a test", "Another test"]






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Any      	=						"Hello I believe in"
         lowercase    :int      	=						pipeline("text-generation"       ,					model="hf-internal-testing/tiny-random-gpt2"				)
         lowercase    :Dict      	=						text_generator(_lowerCAmelCase				)
         self.assertEqual(
             _lowerCAmelCase       ,					[{"generated_text": "Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe"}]       ,					)

         lowercase    :Union[str, Any]      	=						text_generator(_lowerCAmelCase       ,					stop_sequence=" fe"				)
         self.assertEqual(_lowerCAmelCase       ,					[{"generated_text": "Hello I believe in fe"}]				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple       ,					_lowerCAmelCase:					Union[str, Any]       ,					_lowerCAmelCase:					Dict				):
         lowercase    :Any      	=						text_generator.model
         lowercase    :Any      	=						text_generator.tokenizer

         lowercase    :Optional[Any]      	=						text_generator("This is a test"				)
         self.assertEqual(_lowerCAmelCase       ,					[{"generated_text": ANY(_lowerCAmelCase				)}]				)
         self.assertTrue(outputs[0]["generated_text"].startswith("This is a test"				)				)

         lowercase    :Optional[int]      	=						text_generator("This is a test"       ,					return_full_text=_lowerCAmelCase				)
         self.assertEqual(_lowerCAmelCase       ,					[{"generated_text": ANY(_lowerCAmelCase				)}]				)
         self.assertNotIn("This is a test"       ,					outputs[0]["generated_text"]				)

         lowercase    :Optional[Any]      	=						pipeline(task="text-generation"       ,					model=_lowerCAmelCase       ,					tokenizer=_lowerCAmelCase       ,					return_full_text=_lowerCAmelCase				)
         lowercase    :int      	=						text_generator("This is a test"				)
         self.assertEqual(_lowerCAmelCase       ,					[{"generated_text": ANY(_lowerCAmelCase				)}]				)
         self.assertNotIn("This is a test"       ,					outputs[0]["generated_text"]				)

         lowercase    :Optional[int]      	=						text_generator("This is a test"       ,					return_full_text=_lowerCAmelCase				)
         self.assertEqual(_lowerCAmelCase       ,					[{"generated_text": ANY(_lowerCAmelCase				)}]				)
         self.assertTrue(outputs[0]["generated_text"].startswith("This is a test"				)				)

         lowercase    :Dict      	=						text_generator(["This is great !", "Something else"]       ,					num_return_sequences=2       ,					do_sample=_lowerCAmelCase				)
         self.assertEqual(
             _lowerCAmelCase       ,					[
                 [{"generated_text": ANY(_lowerCAmelCase				)}, {"generated_text": ANY(_lowerCAmelCase				)}],
                 [{"generated_text": ANY(_lowerCAmelCase				)}, {"generated_text": ANY(_lowerCAmelCase				)}],
             ]       ,					)

         if text_generator.tokenizer.pad_token is not None:
           lowercase    :List[Any]      	=						text_generator(
               ["This is great !", "Something else"]       ,					num_return_sequences=2       ,					batch_size=2       ,					do_sample=_lowerCAmelCase				)
           self.assertEqual(
               _lowerCAmelCase       ,					[
                   [{"generated_text": ANY(_lowerCAmelCase				)}, {"generated_text": ANY(_lowerCAmelCase				)}],
                   [{"generated_text": ANY(_lowerCAmelCase				)}, {"generated_text": ANY(_lowerCAmelCase				)}],
               ]       ,					)

         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Optional[Any]      	=						text_generator("test"       ,					return_full_text=_lowerCAmelCase       ,					return_text=_lowerCAmelCase				)
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Optional[Any]      	=						text_generator("test"       ,					return_full_text=_lowerCAmelCase       ,					return_tensors=_lowerCAmelCase				)
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Union[str, Any]      	=						text_generator("test"       ,					return_text=_lowerCAmelCase       ,					return_tensors=_lowerCAmelCase				)

         # Empty prompt is slighly special
         # it requires BOS token to exist.
         # Special case for Pegasus which will always append EOS so will
         # work even without BOS.
         if (
             text_generator.tokenizer.bos_token_id is not None
             or "Pegasus" in tokenizer.__class__.__name__
             or "Git" in model.__class__.__name__
         ):
           lowercase    :Optional[Any]      	=						text_generator(""				)
           self.assertEqual(_lowerCAmelCase       ,					[{"generated_text": ANY(_lowerCAmelCase				)}]				)
         else:
           with self.assertRaises((ValueError, AssertionError)				):
             lowercase    :Dict      	=						text_generator(""				)

         if text_generator.framework == "tf":
           # TF generation does not support max_new_tokens, and it's impossible
           # to control long generation with only max_length without
           # fancy calculation, dismissing tests for now.
           return
         # We don't care about infinite range models.
         # They already work.
         # Skip this test for XGLM, since it uses sinusoidal positional embeddings which are resized on-the-fly.
         lowercase    :Any      	=						["RwkvForCausalLM", "XGLMForCausalLM", "GPTNeoXForCausalLM"]
         if (
             tokenizer.model_max_length < 1_00_00
             and text_generator.model.__class__.__name__ not in EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS
         ):
           # Handling of large generations
           with self.assertRaises((RuntimeError, IndexError, ValueError, AssertionError)				):
             text_generator("This is a test" * 5_00       ,					max_new_tokens=20				)

           lowercase    :Optional[int]      	=						text_generator("This is a test" * 5_00       ,					handle_long_generation="hole"       ,					max_new_tokens=20				)
           # Hole strategy cannot work
           with self.assertRaises(_lowerCAmelCase				):
             text_generator(
                 "This is a test" * 5_00       ,					handle_long_generation="hole"       ,					max_new_tokens=tokenizer.model_max_length + 10       ,					)






       @require_torch
       @require_accelerate
       @require_torch_gpu
       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         import torch

         # Classic `model_kwargs`
         lowercase    :Dict      	=						pipeline(
             model="hf-internal-testing/tiny-random-bloom"       ,					model_kwargs={"device_map": "auto", "torch_dtype": torch.bfloataa}       ,					)
         self.assertEqual(pipe.model.device       ,					torch.device(0				)				)
         self.assertEqual(pipe.model.lm_head.weight.dtype       ,					torch.bfloataa				)
         lowercase    :int      	=						pipe("This is a test"				)
         self.assertEqual(
             _lowerCAmelCase       ,					[
                 {
                     "generated_text": (
                         "This is a test test test test test test test test test test test test test test test test"
                         " test"
                     )
                 }
             ]       ,					)

         # Upgraded those two to real pipeline arguments (they just get sent for the model as they're unlikely to mean anything else.)
         lowercase    :Tuple      	=						pipeline(model="hf-internal-testing/tiny-random-bloom"       ,					device_map="auto"       ,					torch_dtype=torch.bfloataa				)
         self.assertEqual(pipe.model.device       ,					torch.device(0				)				)
         self.assertEqual(pipe.model.lm_head.weight.dtype       ,					torch.bfloataa				)
         lowercase    :Optional[Any]      	=						pipe("This is a test"				)
         self.assertEqual(
             _lowerCAmelCase       ,					[
                 {
                     "generated_text": (
                         "This is a test test test test test test test test test test test test test test test test"
                         " test"
                     )
                 }
             ]       ,					)

         # torch_dtype will be automatically set to float32 if not provided - check: https://github.com/huggingface/transformers/pull/20602
         lowercase    :int      	=						pipeline(model="hf-internal-testing/tiny-random-bloom"       ,					device_map="auto"				)
         self.assertEqual(pipe.model.device       ,					torch.device(0				)				)
         self.assertEqual(pipe.model.lm_head.weight.dtype       ,					torch.floataa				)
         lowercase    :List[Any]      	=						pipe("This is a test"				)
         self.assertEqual(
             _lowerCAmelCase       ,					[
                 {
                     "generated_text": (
                         "This is a test test test test test test test test test test test test test test test test"
                         " test"
                     )
                 }
             ]       ,					)






       @require_torch
       @require_torch_gpu
       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         import torch

         lowercase    :List[Any]      	=						pipeline(model="hf-internal-testing/tiny-random-bloom"       ,					device=0       ,					torch_dtype=torch.floataa				)
         pipe("This is a test"				)






       @require_torch
       @require_accelerate
       @require_torch_gpu
       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
         import torch

         lowercase    :str      	=						pipeline(model="hf-internal-testing/tiny-random-bloom"       ,					device_map="auto"       ,					torch_dtype=torch.floataa				)
         pipe("This is a test"       ,					do_sample=_lowerCAmelCase       ,					top_p=0.5				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Dict				):
         lowercase    :Dict      	=						"Hello world"
         lowercase    :Dict      	=						pipeline("text-generation"       ,					model="hf-internal-testing/tiny-random-gpt2"				)
         if text_generator.model.framework == "tf":
           lowercase    :Tuple      	=						logging.get_logger("transformers.generation.tf_utils"				)
         else:
           lowercase    :List[str]      	=						logging.get_logger("transformers.generation.utils"				)
         lowercase    :Optional[Any]      	=						"Both `max_new_tokens`"  # The beggining of the message to be checked in this test

         # Both are set by the user -> log warning
         with CaptureLogger(_lowerCAmelCase				) as cl:
           lowercase    :Union[str, Any]      	=						text_generator(_lowerCAmelCase       ,					max_length=10       ,					max_new_tokens=1				)
         self.assertIn(_lowerCAmelCase       ,					cl.out				)

         # The user only sets one -> no warning
         with CaptureLogger(_lowerCAmelCase				) as cl:
           lowercase    :Optional[int]      	=						text_generator(_lowerCAmelCase       ,					max_new_tokens=1				)
         self.assertNotIn(_lowerCAmelCase       ,					cl.out				)

         with CaptureLogger(_lowerCAmelCase				) as cl:
           lowercase    :Any      	=						text_generator(_lowerCAmelCase       ,					max_length=10				)
         self.assertNotIn(_lowerCAmelCase       ,					cl.out				)









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       UpperCAmelCase__  (     ):
  lowercase    :List[str]      	=						torch.nn.Linear(2,			4   )
  lowercase    :List[Any]      	=						torch.optim.AdamW(model.parameters(),			lr=1.0   )
  lowercase    :int      	=						torch.optim.lr_scheduler.OneCycleLR(lowerCamelCase,			max_lr=0.01,			steps_per_epoch=2,			epochs=1   )
  lowercase    :Optional[int]      	=						DataLoader(TensorDataset(torch.tensor([1, 2, 3]   )   )   )
  lowercase    :List[str]      	=						DataLoader(TensorDataset(torch.tensor([4, 5, 6]   )   )   )

  return model, optimizer, scheduler, train_dl, valid_dl



def       UpperCAmelCase__  (     lowerCamelCase   ):
  return (model.weight.abs().sum() + model.bias.abs().sum()).item()



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :Dict      	=						torch.nn.Linear(*tuple(model.weight.T.shape   )   ).state_dict()
  model.load_state_dict(lowerCamelCase   )







class    __lowerCAmelCase       (							lowerCAmelCase):






       @require_cuda
       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         lowercase    :List[str]      	=						Accelerator()
         assert PartialState._shared_state["_cpu"] is False
         assert PartialState._shared_state["device"].type == "cuda"
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Any      	=						Accelerator(cpu=_lowerCAmelCase				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         lowercase    :Any      	=						Accelerator()
         lowercase    :Dict      	=						GradientState()
         assert state.num_steps == 1
         lowercase    :Tuple      	=						4
         assert state.num_steps == 4

         assert state.sync_gradients is True
         lowercase    :List[Any]      	=						False
         assert state.sync_gradients is False
         GradientState._reset_state()






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :str      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :Any      	=						create_components()

         (
             (
             lowercase
         )	,    (
             lowercase
         )	,    (
             lowercase
         )	,    (
             lowercase
         )	,    (
             lowercase
         )	,    
         )    :Tuple      	=						accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)

         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 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Dict      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :List[str]      	=						create_components()
         accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)
         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 					SCREAMING_SNAKE_CASE    (   self:					int				):
         PartialState._reset_state()

         # Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
         def noop(*_lowerCAmelCase:					List[str]       ,					**_lowerCAmelCase:					Optional[int]				):
           pass

         with patch("torch.cuda.set_device"       ,					_lowerCAmelCase				), patch_environment(ACCELERATE_TORCH_DEVICE="cuda:64"				):
           lowercase    :List[Any]      	=						Accelerator()
           self.assertEqual(str(accelerator.state.device				)       ,					"cuda:64"				)






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :Tuple      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :Optional[Any]      	=						create_components()
         accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)

         lowercase    :Tuple      	=						get_signature(_lowerCAmelCase				)

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

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

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






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         lowercase    :Dict      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :int      	=						create_components()
         accelerator.prepare(_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)

         lowercase    :int      	=						get_signature(_lowerCAmelCase				)

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

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

        # loading hook
         def load_config(_lowerCAmelCase:					Union[str, Any]       ,					_lowerCAmelCase:					Union[str, Any]				):
           with open(os.path.join(_lowerCAmelCase       ,					"data.json"				)       ,					"r"				) as f:
             lowercase    :int      	=						json.load(_lowerCAmelCase				)

           lowercase    :Optional[int]      	=						config["class_name"]

         lowercase    :Optional[Any]      	=						accelerator.register_save_state_pre_hook(_lowerCAmelCase				)
         lowercase    :Tuple      	=						accelerator.register_load_state_pre_hook(_lowerCAmelCase				)

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

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

           # random class name to verify correct one is loaded
           lowercase    :Optional[int]      	=						"random"

           # make sure loaded weights match with hooks
           accelerator.load_state(_lowerCAmelCase				)
           self.assertTrue(abs(model_signature - get_signature(_lowerCAmelCase				)				) < 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(_lowerCAmelCase				)

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

           # random class name to verify correct one is loaded
           lowercase    :str      	=						"random"

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

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






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         lowercase    :List[str]      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :List[Any]      	=						create_components()
         lowercase    :List[str]      	=						None

         # This should work
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :str      	=						accelerator.prepare(
             _lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)
         self.assertTrue(dummy_obj is None				)






       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         lowercase    :Union[str, Any]      	=						Accelerator()
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :Union[str, Any]      	=						create_components()
         lowercase    :Union[str, Any]      	=						[1, 2, 3]

         # This should work
         lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase	,    lowercase    :List[str]      	=						accelerator.prepare(
             _lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase       ,					_lowerCAmelCase				)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Dummy object should have `_is_accelerate_prepared` set to `True`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Model is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Optimizer is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Scheduler is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`"       ,					)
         self.assertEqual(
             getattr(_lowerCAmelCase       ,					"_is_accelerate_prepared"       ,					_lowerCAmelCase				)       ,					_lowerCAmelCase       ,					"Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`"       ,					)






       @slow
       @require_bnb
       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         from transformers import AutoModelForCausalLM

         lowercase    :Any      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					load_in_abit=_lowerCAmelCase       ,					device_map={"": 0}       ,					)
         lowercase    :int      	=						Accelerator()

         # This should work
         lowercase    :Optional[int]      	=						accelerator.prepare(_lowerCAmelCase				)






       @slow
       @require_bnb
       def 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         from transformers import AutoModelForCausalLM

         lowercase    :Optional[Any]      	=						Accelerator()

         with init_empty_weights():
           lowercase    :Any      	=						AutoModelForCausalLM.from_pretrained(
               "EleutherAI/gpt-neo-125m"       ,					)
           model.tie_weights()
           lowercase    :Dict      	=						infer_auto_device_map(_lowerCAmelCase				)
           lowercase    :int      	=						"cpu"

         lowercase    :str      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					device_map=_lowerCAmelCase       ,					load_in_abit=_lowerCAmelCase       ,					llm_inta_enable_fpaa_cpu_offload=_lowerCAmelCase				)

         # This should not work and get value error
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Dict      	=						accelerator.prepare(_lowerCAmelCase				)






       @slow
       @require_bnb
       @require_multi_gpu
       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         from transformers import AutoModelForCausalLM

         lowercase    :List[str]      	=						{"distributed_type": DistributedType.MULTI_GPU}

         with init_empty_weights():
           lowercase    :str      	=						AutoModelForCausalLM.from_pretrained(
               "EleutherAI/gpt-neo-125m"       ,					)
           model.tie_weights()
           lowercase    :Optional[Any]      	=						infer_auto_device_map(_lowerCAmelCase				)
           lowercase    :Any      	=						1

         lowercase    :str      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					load_in_abit=_lowerCAmelCase       ,					device_map=_lowerCAmelCase       ,					)
         lowercase    :List[str]      	=						Accelerator()

         # This should not work and get value error
         with self.assertRaises(_lowerCAmelCase				):
           lowercase    :Optional[Any]      	=						accelerator.prepare(_lowerCAmelCase				)

         PartialState._reset_state()






       @slow
       @require_bnb
       @require_multi_gpu
       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
         from transformers import AutoModelForCausalLM

         with init_empty_weights():
           lowercase    :List[str]      	=						AutoModelForCausalLM.from_pretrained(
               "EleutherAI/gpt-neo-125m"       ,					)
           lowercase    :List[Any]      	=						infer_auto_device_map(_lowerCAmelCase				)
           lowercase    :Optional[int]      	=						1

         lowercase    :Union[str, Any]      	=						AutoModelForCausalLM.from_pretrained(
             "EleutherAI/gpt-neo-125m"       ,					load_in_abit=_lowerCAmelCase       ,					device_map=_lowerCAmelCase       ,					)
         lowercase    :Union[str, Any]      	=						Accelerator()

         # This should work
         lowercase    :List[Any]      	=						accelerator.prepare(_lowerCAmelCase				)






       @require_cuda
       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Optional[int]      	=						torch.nn.Linear(10       ,					10				)
         lowercase    :Optional[int]      	=						torch.optim.SGD(model.parameters()       ,					lr=0.01				)
         lowercase    :List[Any]      	=						Accelerator(cpu=_lowerCAmelCase				)
         lowercase    :List[str]      	=						accelerator.prepare(_lowerCAmelCase				)



