




'''simple docstring'''






import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


_lowerCAmelCase       =   [nn.LayerNorm]

_lowerCAmelCase       =   logging.get_logger(__name__)

_lowerCAmelCase       =   version.parse(version.parse(torch.__version__).base_version)

_lowerCAmelCase       =   parsed_torch_version_base >= version.parse('''2.0''')
_lowerCAmelCase       =   parsed_torch_version_base >= version.parse('''1.12''')
_lowerCAmelCase       =   parsed_torch_version_base >= version.parse('''1.11''')
_lowerCAmelCase       =   parsed_torch_version_base < version.parse('''1.11''')

def      __lowerCAmelCase   (       snake_case__   , snake_case__   , snake_case__   , snake_case__   , self		):
			from torch import _softmax_backward_data

			if is_torch_less_than_1_11:
						return _softmax_backward_data(snake_case__   , snake_case__   , parent.dim   , self		)
			else:
						return _softmax_backward_data(snake_case__   , snake_case__   , parent.dim   , self.dtype		)

def      __lowerCAmelCase   (       snake_case__   , snake_case__   , snake_case__ = 0		):
			__UpperCamelCase :     List[Any]										=  index.to(layer.weight.device		)
			__UpperCamelCase :     Optional[int]										=  layer.weight.index_select(snake_case__   , snake_case__		).clone().detach()
			if layer.bias is not None:
						if dim == 1:
									__UpperCamelCase :     List[Any]										=  layer.bias.clone().detach()
						else:
									__UpperCamelCase :     List[str]										=  layer.bias[index].clone().detach()
			__UpperCamelCase :     Union[str, Any]										=  list(layer.weight.size()		)
			__UpperCamelCase :     Any										=  len(snake_case__		)
			__UpperCamelCase :     List[str]										=  nn.Linear(new_size[1]   , new_size[0]   , bias=layer.bias is not None		).to(layer.weight.device		)
			__UpperCamelCase :     str										=  False
			new_layer.weight.copy_(W.contiguous()		)
			__UpperCamelCase :     List[Any]										=  True
			if layer.bias is not None:
						__UpperCamelCase :     int										=  False
						new_layer.bias.copy_(b.contiguous()		)
						__UpperCamelCase :     Any										=  True
			return new_layer



class    A						(					nn.Module						):




				'''simple docstring'''





				def __init__(self					,    _UpperCAmelCase					,    _UpperCAmelCase     )      ->			Tuple:
							super().__init__()
							__UpperCamelCase :     Any										=  nf
							__UpperCamelCase :     List[str]										=  nn.Parameter(torch.empty(_UpperCAmelCase					,    _UpperCAmelCase     )     )
							__UpperCamelCase :     Any										=  nn.Parameter(torch.zeros(_UpperCAmelCase     )     )
							nn.init.normal_(self.weight					,    std=0.02     )





				def 			a_			(self					,    _UpperCAmelCase     )      ->			Optional[int]:
							__UpperCamelCase :     Dict										=  x.size()[:-1] + (self.nf,)
							__UpperCamelCase :     List[str]										=  torch.addmm(self.bias					,    x.view(-1					,    x.size(-1     )     )					,    self.weight     )
							__UpperCamelCase :     Optional[Any]										=  x.view(_UpperCAmelCase     )
							return x







def      __lowerCAmelCase   (       snake_case__   , snake_case__   , snake_case__ = 1		):
			__UpperCamelCase :     Union[str, Any]										=  index.to(layer.weight.device		)
			__UpperCamelCase :     str										=  layer.weight.index_select(snake_case__   , snake_case__		).clone().detach()
			if dim == 0:
						__UpperCamelCase :     int										=  layer.bias.clone().detach()
			else:
						__UpperCamelCase :     Optional[int]										=  layer.bias[index].clone().detach()
			__UpperCamelCase :     Dict										=  list(layer.weight.size()		)
			__UpperCamelCase :     List[Any]										=  len(snake_case__		)
			__UpperCamelCase :     List[Any]										=  ConvaD(new_size[1]   , new_size[0]		).to(layer.weight.device		)
			__UpperCamelCase :     Optional[Any]										=  False
			new_layer.weight.copy_(W.contiguous()		)
			__UpperCamelCase :     Union[str, Any]										=  True
			__UpperCamelCase :     List[str]										=  False
			new_layer.bias.copy_(b.contiguous()		)
			__UpperCamelCase :     int										=  True
			return new_layer

def      __lowerCAmelCase   (       snake_case__   , snake_case__   , snake_case__ = None		):
			if isinstance(snake_case__   , nn.Linear		):
						return prune_linear_layer(snake_case__   , snake_case__   , dim=0 if dim is None else dim		)
			elif isinstance(snake_case__   , snake_case__		):
						return prune_convad_layer(snake_case__   , snake_case__   , dim=1 if dim is None else dim		)
			else:
						raise ValueError(F"Can't prune layer of class {layer.__class__}"		)

def      __lowerCAmelCase   (       snake_case__   , snake_case__   , snake_case__   , *snake_case__		):
			assert len(snake_case__		) > 0, F"{input_tensors} has to be a tuple/list of tensors"

			# inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
			__UpperCamelCase :     Any										=  len(inspect.signature(snake_case__		).parameters		)
			if num_args_in_forward_chunk_fn != len(snake_case__		):
						raise ValueError(
						    F"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(snake_case__		)} input "
						    "tensors are given"		)

			if chunk_size > 0:
						__UpperCamelCase :     Any										=  input_tensors[0].shape[chunk_dim]
						for input_tensor in input_tensors:
									if input_tensor.shape[chunk_dim] != tensor_shape:
												raise ValueError(
												    F"All input tenors have to be of the same shape: {tensor_shape}, "
												    F"found shape {input_tensor.shape[chunk_dim]}"		)

						if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
									raise ValueError(
									    F"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
									    F"size {chunk_size}"		)

						__UpperCamelCase :     Optional[Any]										=  input_tensors[0].shape[chunk_dim] // chunk_size

						# chunk input tensor into tuples
						__UpperCamelCase :     Tuple										=  tuple(input_tensor.chunk(snake_case__   , dim=snake_case__		) for input_tensor in input_tensors		)
						# apply forward fn to every tuple
						__UpperCamelCase :     Any										=  tuple(forward_fn(*snake_case__		) for input_tensors_chunk in zip(*snake_case__		)		)
						# concatenate output at same dimension
						return torch.cat(snake_case__   , dim=snake_case__		)

			return forward_fn(*snake_case__		)

def      __lowerCAmelCase   (       snake_case__   , snake_case__   , snake_case__   , snake_case__		):
			__UpperCamelCase :     List[str]										=  torch.ones(snake_case__   , snake_case__		)
			__UpperCamelCase :     Union[str, Any]										=  set(snake_case__		) - already_pruned_heads  # Convert to set and remove already pruned heads
			for head in heads:
						# Compute how many pruned heads are before the head and move the index accordingly
						__UpperCamelCase :     Tuple										=  head - sum(1 if h < head else 0 for h in already_pruned_heads		)
						__UpperCamelCase :     Optional[int]										=  0
			__UpperCamelCase :     List[str]										=  mask.view(-1		).contiguous().eq(1		)
			__UpperCamelCase :     torch.LongTensor										=  torch.arange(len(snake_case__		)		)[mask].long()
			return heads, index

def      __lowerCAmelCase   (       *snake_case__   , snake_case__ = None		):
			return torch.meshgrid(*snake_case__   , indexing=snake_case__		)




def      __lowerCAmelCase   (       snake_case__		):
			return tensor.device, storage_ptr(snake_case__		), storage_size(snake_case__		)









'''simple docstring'''






import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


_lowerCAmelCase       =   namedtuple('''ExampleNamedTuple''', '''a b c''')



class    A						(					unittest.TestCase						):




				'''simple docstring'''





				def 			a_			(self     )      ->			Tuple:
							__UpperCamelCase :     Optional[int]										=  torch.randn(5					,    2     )
							__UpperCamelCase :     List[str]										=  torch.device("cuda"     ) if torch.cuda.is_available() else torch.device("cpu"     )

							__UpperCamelCase :     Optional[Any]										=  send_to_device(_UpperCAmelCase					,    _UpperCAmelCase     )
							self.assertTrue(torch.equal(resulta.cpu()					,    _UpperCAmelCase     )     )

							__UpperCamelCase :     Any										=  send_to_device((tensor, [tensor, tensor], 1)					,    _UpperCAmelCase     )
							self.assertIsInstance(_UpperCAmelCase					,    _UpperCAmelCase     )
							self.assertTrue(torch.equal(resulta[0].cpu()					,    _UpperCAmelCase     )     )
							self.assertIsInstance(resulta[1]					,    _UpperCAmelCase     )
							self.assertTrue(torch.equal(resulta[1][0].cpu()					,    _UpperCAmelCase     )     )
							self.assertTrue(torch.equal(resulta[1][1].cpu()					,    _UpperCAmelCase     )     )
							self.assertEqual(resulta[2]					,    1     )

							__UpperCamelCase :     Optional[int]										=  send_to_device({"a": tensor, "b": [tensor, tensor], "c": 1}					,    _UpperCAmelCase     )
							self.assertIsInstance(_UpperCAmelCase					,    _UpperCAmelCase     )
							self.assertTrue(torch.equal(resulta["a"].cpu()					,    _UpperCAmelCase     )     )
							self.assertIsInstance(resulta["b"]					,    _UpperCAmelCase     )
							self.assertTrue(torch.equal(resulta["b"][0].cpu()					,    _UpperCAmelCase     )     )
							self.assertTrue(torch.equal(resulta["b"][1].cpu()					,    _UpperCAmelCase     )     )
							self.assertEqual(resulta["c"]					,    1     )

							__UpperCamelCase :     Optional[Any]										=  send_to_device(ExampleNamedTuple(a=_UpperCAmelCase					,    b=[tensor, tensor]					,    c=1     )					,    _UpperCAmelCase     )
							self.assertIsInstance(_UpperCAmelCase					,    _UpperCAmelCase     )
							self.assertTrue(torch.equal(resulta.a.cpu()					,    _UpperCAmelCase     )     )
							self.assertIsInstance(resulta.b					,    _UpperCAmelCase     )
							self.assertTrue(torch.equal(resulta.b[0].cpu()					,    _UpperCAmelCase     )     )
							self.assertTrue(torch.equal(resulta.b[1].cpu()					,    _UpperCAmelCase     )     )
							self.assertEqual(resulta.c					,    1     )

							__UpperCamelCase :     Dict										=  send_to_device(UserDict({"a": tensor, "b": [tensor, tensor], "c": 1}     )					,    _UpperCAmelCase     )
							self.assertIsInstance(_UpperCAmelCase					,    _UpperCAmelCase     )
							self.assertTrue(torch.equal(resulta["a"].cpu()					,    _UpperCAmelCase     )     )
							self.assertIsInstance(resulta["b"]					,    _UpperCAmelCase     )
							self.assertTrue(torch.equal(resulta["b"][0].cpu()					,    _UpperCAmelCase     )     )
							self.assertTrue(torch.equal(resulta["b"][1].cpu()					,    _UpperCAmelCase     )     )
							self.assertEqual(resulta["c"]					,    1     )
				def 			a_			(self     )      ->			Optional[Any]:
							with self.assertRaises(_UpperCAmelCase     ) as cm:
										__UpperCamelCase :     int										=  recursively_apply(torch.tensor					,    (torch.tensor(1     ), 1)					,    error_on_other_type=_UpperCAmelCase     )
							self.assertEqual(
							    str(cm.exception     )					,    "Unsupported types (<class 'int'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed."					,    )
				def 			a_			(self     )      ->			Union[str, Any]:
							__UpperCamelCase :     Dict										=  torch.tensor([1, 2, 3, 4, 5]     )
							self.assertEqual(listify(_UpperCAmelCase     )					,    [1, 2, 3, 4, 5]     )

							__UpperCamelCase :     int										=  torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]     )
							self.assertEqual(listify(_UpperCAmelCase     )					,    [[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]     )

							__UpperCamelCase :     Dict										=  torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]     )
							self.assertEqual(
							    listify(_UpperCAmelCase     )					,    [[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]     )
				def 			a_			(self     )      ->			Optional[Any]:
							with patch_environment(aa=1					,    BB=2     ):
										self.assertEqual(os.environ.get("AA"     )					,    "1"     )
										self.assertEqual(os.environ.get("BB"     )					,    "2"     )

							self.assertNotIn("AA"					,    os.environ     )
							self.assertNotIn("BB"					,    os.environ     )
				def 			a_			(self     )      ->			Dict:
							__UpperCamelCase :     Tuple										=  RegressionModel()
							__UpperCamelCase :     List[str]										=  model.forward
							__UpperCamelCase :     Any										=  convert_outputs_to_fpaa(model.forward     )
							__UpperCamelCase :     List[str]										=  extract_model_from_parallel(_UpperCAmelCase					,    keep_fpaa_wrapper=_UpperCAmelCase     )
							__UpperCamelCase :     Dict										=  pickle.dumps(_UpperCAmelCase     )
				@require_cuda
				def 			a_			(self     )      ->			str:
							__UpperCamelCase :     List[str]										=  RegressionModel()
							__UpperCamelCase :     List[Any]										=  model.forward
							__UpperCamelCase :     List[str]										=  torch.cuda.amp.autocast(dtype=torch.floataa     )(model.forward     )
							__UpperCamelCase :     str										=  convert_outputs_to_fpaa(model.forward     )
							__UpperCamelCase :     int										=  extract_model_from_parallel(_UpperCAmelCase					,    keep_fpaa_wrapper=_UpperCAmelCase     )
							__UpperCamelCase :     int										=  pickle.dumps(_UpperCAmelCase     )
				@require_cuda
				@require_torch_min_version(version="2.0"     )
				def 			a_			(self     )      ->			Tuple:
							__UpperCamelCase :     List[Any]										=  RegressionModel()
							__UpperCamelCase :     Tuple										=  model.forward
							__UpperCamelCase :     Optional[int]										=  torch.cuda.amp.autocast(dtype=torch.floataa     )(model.forward     )
							__UpperCamelCase :     List[Any]										=  convert_outputs_to_fpaa(model.forward     )
							__UpperCamelCase :     Union[str, Any]										=  torch.compile(model.forward					,    backend="inductor"     )
							__UpperCamelCase :     int										=  torch.randn(4					,    1_0     ).cuda()
							__UpperCamelCase :     int										=  model(_UpperCAmelCase     )
				def 			a_			(self     )      ->			Union[str, Any]:
							__UpperCamelCase :     List[Any]										=  RegressionModel()
							# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
							__UpperCamelCase :     Optional[Any]										=  torch.nn.parallel.DataParallel(_UpperCAmelCase     )
							__UpperCamelCase :     Union[str, Any]										=  extract_model_from_parallel(_UpperCAmelCase     )

							self.assertEqual(_UpperCAmelCase					,    _UpperCAmelCase     )
				@require_torch_min_version(version="2.0"     )
				def 			a_			(self     )      ->			Union[str, Any]:
							__UpperCamelCase :     int										=  RegressionModel()
							__UpperCamelCase :     List[str]										=  torch.compile(_UpperCAmelCase     )

							# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
							__UpperCamelCase :     Union[str, Any]										=  torch.nn.parallel.DataParallel(_UpperCAmelCase     )
							__UpperCamelCase :     List[str]										=  torch.compile(_UpperCAmelCase     )
							__UpperCamelCase :     Optional[int]										=  extract_model_from_parallel(_UpperCAmelCase     )

							self.assertEqual(compiled_model._orig_mod					,    compiled_model_unwrapped._orig_mod     )





				def 			a_			(self     )      ->			str:
							self.assertEqual(find_device([1, "a", torch.tensor([1, 2, 3]     )]     )					,    torch.device("cpu"     )     )
							self.assertEqual(find_device({"a": 1, "b": torch.tensor([1, 2, 3]     )}     )					,    torch.device("cpu"     )     )
							self.assertIsNone(find_device([1, "a"]     )     )



