'''simple docstring'''




import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu






class 							a__(   nn.Module      ):


			'''simple docstring'''






			def __init__(  self):




					"""simple docstring"""


					super().__init__()
					lowerCAmelCase     =    nn.Linear(3			,       4)
					lowerCAmelCase     =    nn.BatchNormad(4)
					lowerCAmelCase     =    nn.Linear(4			,       5)







			def 				a_			(  self			,       __lowerCAmelCase):




					"""simple docstring"""


					return self.lineara(self.batchnorm(self.lineara(__lowerCAmelCase)))







class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			def 				a_			(  self			,       __lowerCAmelCase			,       *__lowerCAmelCase			,       **__lowerCAmelCase):




					"""simple docstring"""


					return (args[0] + 1,) + args[1:], kwargs







class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase):




					"""simple docstring"""


					return output + 1







class 							a__(   unittest.TestCase      ):


			'''simple docstring'''






			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ModelForTest()
					lowerCAmelCase     =    ModelHook()

					add_hook_to_module(__lowerCAmelCase			,       __lowerCAmelCase)
					self.assertEqual(test_model._hf_hook			,       __lowerCAmelCase)
					self.assertTrue(hasattr(__lowerCAmelCase			,       """_old_forward"""))

					# Check adding the hook did not change the name or the signature
					self.assertEqual(test_model.forward.__name__			,       """forward""")
					self.assertListEqual(list(inspect.signature(test_model.forward).parameters)			,       ["""x"""])

					remove_hook_from_module(__lowerCAmelCase)
					self.assertFalse(hasattr(__lowerCAmelCase			,       """_hf_hook"""))
					self.assertFalse(hasattr(__lowerCAmelCase			,       """_old_forward"""))



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ModelForTest()
					lowerCAmelCase     =    ModelHook()

					add_hook_to_module(__lowerCAmelCase			,       __lowerCAmelCase)
					add_hook_to_module(__lowerCAmelCase			,       __lowerCAmelCase			,       append=__lowerCAmelCase)

					self.assertEqual(isinstance(test_model._hf_hook			,       __lowerCAmelCase)			,       __lowerCAmelCase)
					self.assertEqual(len(test_model._hf_hook.hooks)			,       2)
					self.assertTrue(hasattr(__lowerCAmelCase			,       """_old_forward"""))

					# Check adding the hook did not change the name or the signature
					self.assertEqual(test_model.forward.__name__			,       """forward""")
					self.assertListEqual(list(inspect.signature(test_model.forward).parameters)			,       ["""x"""])

					remove_hook_from_module(__lowerCAmelCase)
					self.assertFalse(hasattr(__lowerCAmelCase			,       """_hf_hook"""))
					self.assertFalse(hasattr(__lowerCAmelCase			,       """_old_forward"""))



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ModelForTest()
					lowerCAmelCase     =    torch.randn(2			,       3)
					lowerCAmelCase     =    test_model(x + 1)
					lowerCAmelCase     =    test_model(x + 2)

					lowerCAmelCase     =    PreForwardHook()
					add_hook_to_module(__lowerCAmelCase			,       __lowerCAmelCase)
					lowerCAmelCase     =    test_model(__lowerCAmelCase)
					self.assertTrue(torch.allclose(__lowerCAmelCase			,       __lowerCAmelCase			,       atol=1E-5))

					# Attaching a hook to a model when it already has one replaces, does not chain
					lowerCAmelCase     =    PreForwardHook()
					add_hook_to_module(__lowerCAmelCase			,       __lowerCAmelCase)
					lowerCAmelCase     =    test_model(__lowerCAmelCase)
					self.assertTrue(torch.allclose(__lowerCAmelCase			,       __lowerCAmelCase			,       atol=1E-5))

					# You need to use the sequential hook to chain two or more hooks
					lowerCAmelCase     =    SequentialHook(PreForwardHook()			,       PreForwardHook())
					add_hook_to_module(__lowerCAmelCase			,       __lowerCAmelCase)

					lowerCAmelCase     =    test_model(__lowerCAmelCase)
					assert torch.allclose(__lowerCAmelCase			,       __lowerCAmelCase			,       atol=1E-5)



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ModelForTest()
					lowerCAmelCase     =    torch.randn(2			,       3)
					lowerCAmelCase     =    test_model(__lowerCAmelCase)

					lowerCAmelCase     =    PostForwardHook()
					add_hook_to_module(__lowerCAmelCase			,       __lowerCAmelCase)
					lowerCAmelCase     =    test_model(__lowerCAmelCase)
					self.assertTrue(torch.allclose(__lowerCAmelCase			,       output + 1			,       atol=1E-5))

					# Attaching a hook to a model when it already has one replaces, does not chain
					lowerCAmelCase     =    PostForwardHook()
					add_hook_to_module(__lowerCAmelCase			,       __lowerCAmelCase)
					lowerCAmelCase     =    test_model(__lowerCAmelCase)
					self.assertTrue(torch.allclose(__lowerCAmelCase			,       output + 1			,       atol=1E-5))

					# You need to use the sequential hook to chain two or more hooks
					lowerCAmelCase     =    SequentialHook(PostForwardHook()			,       PostForwardHook())
					add_hook_to_module(__lowerCAmelCase			,       __lowerCAmelCase)

					lowerCAmelCase     =    test_model(__lowerCAmelCase)
					assert torch.allclose(__lowerCAmelCase			,       output + 2			,       atol=1E-5)



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ModelForTest()
					lowerCAmelCase     =    torch.randn(2			,       3)
					lowerCAmelCase     =    test_model(__lowerCAmelCase)

					lowerCAmelCase     =    PostForwardHook()
					add_hook_to_module(__lowerCAmelCase			,       __lowerCAmelCase)
					lowerCAmelCase     =    test_model(__lowerCAmelCase)
					self.assertTrue(torch.allclose(__lowerCAmelCase			,       output + 1))
					self.assertTrue(outputa.requires_grad)

					lowerCAmelCase     =    True
					lowerCAmelCase     =    test_model(__lowerCAmelCase)
					self.assertFalse(outputa.requires_grad)



			@require_multi_gpu
			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ModelForTest()
					# Everything is on CPU
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))

					# This will move each submodule on different devices
					add_hook_to_module(model.lineara			,       AlignDevicesHook(execution_device=0))
					add_hook_to_module(model.batchnorm			,       AlignDevicesHook(execution_device=0))
					add_hook_to_module(model.lineara			,       AlignDevicesHook(execution_device=1))

					self.assertEqual(model.lineara.weight.device			,       torch.device(0))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device(0))
					self.assertEqual(model.batchnorm.running_mean.device			,       torch.device(0))
					self.assertEqual(model.lineara.weight.device			,       torch.device(1))

					# We can still make a forward pass. The input does not need to be on any particular device
					lowerCAmelCase     =    torch.randn(2			,       3)
					lowerCAmelCase     =    model(__lowerCAmelCase)
					self.assertEqual(output.device			,       torch.device(1))

					# We can add a general hook to put back output on same device as input.
					add_hook_to_module(__lowerCAmelCase			,       AlignDevicesHook(io_same_device=__lowerCAmelCase))
					lowerCAmelCase     =    torch.randn(2			,       3).to(0)
					lowerCAmelCase     =    model(__lowerCAmelCase)
					self.assertEqual(output.device			,       torch.device(0))



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ModelForTest()

					# Everything is on CPU
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))

					# This will move each submodule on different devices
					lowerCAmelCase     =    {"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

					add_hook_to_module(model.lineara			,       AlignDevicesHook(**__lowerCAmelCase))
					add_hook_to_module(model.batchnorm			,       AlignDevicesHook(**__lowerCAmelCase))
					add_hook_to_module(model.lineara			,       AlignDevicesHook(**__lowerCAmelCase))

					# Parameters have been offloaded, so on the meta device
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					# Buffers are not included in the offload by default, so are on the execution device
					lowerCAmelCase     =    torch.device(hook_kwargs["""execution_device"""])
					self.assertEqual(model.batchnorm.running_mean.device			,       __lowerCAmelCase)

					lowerCAmelCase     =    torch.randn(2			,       3)
					lowerCAmelCase     =    model(__lowerCAmelCase)
					self.assertEqual(output.device			,       __lowerCAmelCase)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_module(model.lineara)
					remove_hook_from_module(model.batchnorm)
					remove_hook_from_module(model.lineara)
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))

					# Now test with buffers included in the offload
					lowerCAmelCase     =    {
					    """execution_device""": 0 if torch.cuda.is_available() else """cpu""",
					    """offload""": True,
					    """offload_buffers""": True,
					}

					add_hook_to_module(model.lineara			,       AlignDevicesHook(**__lowerCAmelCase))
					add_hook_to_module(model.batchnorm			,       AlignDevicesHook(**__lowerCAmelCase))
					add_hook_to_module(model.lineara			,       AlignDevicesHook(**__lowerCAmelCase))

					# Parameters have been offloaded, so on the meta device, buffers included
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.batchnorm.running_mean.device			,       torch.device("""meta"""))

					lowerCAmelCase     =    torch.randn(2			,       3)
					lowerCAmelCase     =    model(__lowerCAmelCase)
					self.assertEqual(output.device			,       __lowerCAmelCase)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_module(model.lineara)
					remove_hook_from_module(model.batchnorm)
					remove_hook_from_module(model.lineara)
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ModelForTest()

					# Everything is on CPU
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))

					# This will move each submodule on different devices
					lowerCAmelCase     =    0 if torch.cuda.is_available() else """cpu"""
					attach_align_device_hook(__lowerCAmelCase			,       execution_device=__lowerCAmelCase			,       offload=__lowerCAmelCase)

					# Parameters have been offloaded, so on the meta device
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					# Buffers are not included in the offload by default, so are on the execution device
					lowerCAmelCase     =    torch.device(__lowerCAmelCase)
					self.assertEqual(model.batchnorm.running_mean.device			,       __lowerCAmelCase)

					lowerCAmelCase     =    torch.randn(2			,       3)
					lowerCAmelCase     =    model(__lowerCAmelCase)
					self.assertEqual(output.device			,       __lowerCAmelCase)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_submodules(__lowerCAmelCase)
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))

					# Now test with buffers included in the offload
					attach_align_device_hook(__lowerCAmelCase			,       execution_device=__lowerCAmelCase			,       offload=__lowerCAmelCase			,       offload_buffers=__lowerCAmelCase)

					# Parameters have been offloaded, so on the meta device, buffers included
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.batchnorm.running_mean.device			,       torch.device("""meta"""))

					lowerCAmelCase     =    torch.randn(2			,       3)
					lowerCAmelCase     =    model(__lowerCAmelCase)
					self.assertEqual(output.device			,       __lowerCAmelCase)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_submodules(__lowerCAmelCase)
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))







			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ModelForTest()

					# Everything is on CPU
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))

					# This will move each submodule on different devices
					lowerCAmelCase     =    0 if torch.cuda.is_available() else """cpu"""
					attach_align_device_hook(
					    __lowerCAmelCase			,       execution_device=__lowerCAmelCase			,       offload=__lowerCAmelCase			,       weights_map=model.state_dict())

					# Parameters have been offloaded, so on the meta device
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					# Buffers are not included in the offload by default, so are on the execution device
					lowerCAmelCase     =    torch.device(__lowerCAmelCase)
					self.assertEqual(model.batchnorm.running_mean.device			,       __lowerCAmelCase)

					lowerCAmelCase     =    torch.randn(2			,       3)
					lowerCAmelCase     =    model(__lowerCAmelCase)
					self.assertEqual(output.device			,       __lowerCAmelCase)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_submodules(__lowerCAmelCase)
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))

					# Now test with buffers included in the offload
					attach_align_device_hook(
					    __lowerCAmelCase			,       execution_device=__lowerCAmelCase			,       offload=__lowerCAmelCase			,       weights_map=model.state_dict()			,       offload_buffers=__lowerCAmelCase			,       )

					# Parameters have been offloaded, so on the meta device, buffers included
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""meta"""))
					self.assertEqual(model.batchnorm.running_mean.device			,       torch.device("""meta"""))

					lowerCAmelCase     =    torch.randn(2			,       3)
					lowerCAmelCase     =    model(__lowerCAmelCase)
					self.assertEqual(output.device			,       __lowerCAmelCase)

					# Removing hooks loads back the weights in the model.
					remove_hook_from_submodules(__lowerCAmelCase)
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.batchnorm.weight.device			,       torch.device("""cpu"""))
					self.assertEqual(model.lineara.weight.device			,       torch.device("""cpu"""))





'''simple docstring'''




import math
import time

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
						import torch_xla.core.xla_model as xm
						import torch_xla.debug.metrics as met






class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			def __init__(  self			,       *__lowerCAmelCase			,       __lowerCAmelCase=None			,       __lowerCAmelCase=None			,       **__lowerCAmelCase):




					"""simple docstring"""


					super().__init__(*__lowerCAmelCase			,       **__lowerCAmelCase)
					lowerCAmelCase     =    eval_examples
					lowerCAmelCase     =    post_process_function



			def 				a_			(  self			,       __lowerCAmelCase=None			,       __lowerCAmelCase=None			,       __lowerCAmelCase=None			,       __lowerCAmelCase = "eval"):




					"""simple docstring"""


					lowerCAmelCase     =    self.eval_dataset if eval_dataset is None else eval_dataset
					lowerCAmelCase     =    self.get_eval_dataloader(__lowerCAmelCase)
					lowerCAmelCase     =    self.eval_examples if eval_examples is None else eval_examples

					# Temporarily disable metric computation, we will do it in the loop here.
					lowerCAmelCase     =    self.compute_metrics
					lowerCAmelCase     =    None
					lowerCAmelCase     =    self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					lowerCAmelCase     =    time.time()
					try:
							lowerCAmelCase     =    eval_loop(
							    __lowerCAmelCase			,       description="""Evaluation"""			,       prediction_loss_only=True if compute_metrics is None else None			,       ignore_keys=__lowerCAmelCase			,       metric_key_prefix=__lowerCAmelCase			,       )
					finally:
							lowerCAmelCase     =    compute_metrics
					lowerCAmelCase     =    self.args.eval_batch_size * self.args.world_size
					if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
							start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
					output.metrics.update(
					    speed_metrics(
					        __lowerCAmelCase			,       __lowerCAmelCase			,       num_samples=output.num_samples			,       num_steps=math.ceil(output.num_samples / total_batch_size)			,       ))
					if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
							# Only the main node write the results by default
							lowerCAmelCase     =    self.post_process_function(__lowerCAmelCase			,       __lowerCAmelCase			,       output.predictions)
							lowerCAmelCase     =    self.compute_metrics(__lowerCAmelCase)

							# Prefix all keys with metric_key_prefix + '_'
							for key in list(metrics.keys()):
									if not key.startswith(f"{metric_key_prefix}_"):
											lowerCAmelCase     =    metrics.pop(__lowerCAmelCase)
							metrics.update(output.metrics)
					else:
							lowerCAmelCase     =    output.metrics

					if self.args.should_log:
							# Only the main node log the results by default
							self.log(__lowerCAmelCase)

					if self.args.tpu_metrics_debug or self.args.debug:
							# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
							xm.master_print(met.metrics_report())

					lowerCAmelCase     =    self.callback_handler.on_evaluate(self.args			,       self.state			,       self.control			,       __lowerCAmelCase)
					return metrics







			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase=None			,       __lowerCAmelCase = "test"):




					"""simple docstring"""


					lowerCAmelCase     =    self.get_test_dataloader(__lowerCAmelCase)

					# Temporarily disable metric computation, we will do it in the loop here.
					lowerCAmelCase     =    self.compute_metrics
					lowerCAmelCase     =    None
					lowerCAmelCase     =    self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					lowerCAmelCase     =    time.time()
					try:
							lowerCAmelCase     =    eval_loop(
							    __lowerCAmelCase			,       description="""Prediction"""			,       prediction_loss_only=True if compute_metrics is None else None			,       ignore_keys=__lowerCAmelCase			,       metric_key_prefix=__lowerCAmelCase			,       )
					finally:
							lowerCAmelCase     =    compute_metrics
					lowerCAmelCase     =    self.args.eval_batch_size * self.args.world_size
					if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
							start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
					output.metrics.update(
					    speed_metrics(
					        __lowerCAmelCase			,       __lowerCAmelCase			,       num_samples=output.num_samples			,       num_steps=math.ceil(output.num_samples / total_batch_size)			,       ))

					if self.post_process_function is None or self.compute_metrics is None:
							return output

					lowerCAmelCase     =    self.post_process_function(__lowerCAmelCase			,       __lowerCAmelCase			,       output.predictions			,       """predict""")
					lowerCAmelCase     =    self.compute_metrics(__lowerCAmelCase)

					# Prefix all keys with metric_key_prefix + '_'
					for key in list(metrics.keys()):
							if not key.startswith(f"{metric_key_prefix}_"):
									lowerCAmelCase     =    metrics.pop(__lowerCAmelCase)
					metrics.update(output.metrics)
					return PredictionOutput(predictions=predictions.predictions			,       label_ids=predictions.label_ids			,       metrics=__lowerCAmelCase)




