




'''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


UpperCamelCase_		:    Union[str, Any]									=		[nn.LayerNorm]

UpperCamelCase_		:    Union[str, Any]									=		logging.get_logger(__name__)

UpperCamelCase_		:    Optional[Any]									=		version.parse(version.parse(torch.__version__).base_version)

UpperCamelCase_		:    Optional[Any]									=		parsed_torch_version_base >= version.parse('''2.0''')
UpperCamelCase_		:    str									=		parsed_torch_version_base >= version.parse('''1.12''')
UpperCamelCase_		:    int									=		parsed_torch_version_base >= version.parse('''1.11''')
UpperCamelCase_		:    int									=		parsed_torch_version_base < version.parse('''1.11''')





def 			__a			(     _UpperCamelCase:					Optional[Any]      ,		_UpperCamelCase:					str      ,		_UpperCamelCase:					Tuple      ,		_UpperCamelCase:					Dict      ,		self:					str					)		->				List[Any]:

     """simple docstring"""




     from torch import _softmax_backward_data

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





def 			__a			(     _UpperCamelCase:					nn.Linear      ,		_UpperCamelCase:					torch.LongTensor      ,		_UpperCamelCase:					int = 0					)		->				nn.Linear:

     """simple docstring"""




     _snake_case							=   index.to(layer.weight.device					)
     _snake_case							=   layer.weight.index_select(_UpperCamelCase      ,		_UpperCamelCase					).clone().detach()
     if layer.bias is not None:
          if dim == 1:
               _snake_case							=   layer.bias.clone().detach()
          else:
               _snake_case							=   layer.bias[index].clone().detach()
     _snake_case							=   list(layer.weight.size()					)
     _snake_case							=   len(_UpperCamelCase					)
     _snake_case							=   nn.Linear(new_size[1]      ,		new_size[0]      ,		bias=layer.bias is not None					).to(layer.weight.device					)
     _snake_case							=   False
     new_layer.weight.copy_(W.contiguous()					)
     _snake_case							=   True
     if layer.bias is not None:
          _snake_case							=   False
          new_layer.bias.copy_(b.contiguous()					)
          _snake_case							=   True
     return new_layer

class      _a       (		nn.Module							):

     '''simple docstring'''





     def __init__(				self      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)							->    List[Any]:
          super().__init__()
          _snake_case							=   nf
          _snake_case							=   nn.Parameter(torch.empty(_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE					)					)
          _snake_case							=   nn.Parameter(torch.zeros(_SCREAMING_SNAKE_CASE					)					)
          nn.init.normal_(self.weight      ,std=0.0_2					)






     def        _lowercase			(				self      ,_SCREAMING_SNAKE_CASE					)							->    int:
          _snake_case							=   x.size()[:-1] + (self.nf,)
          _snake_case							=   torch.addmm(self.bias      ,x.view(-1      ,x.size(-1					)					)      ,self.weight					)
          _snake_case							=   x.view(_SCREAMING_SNAKE_CASE					)
          return x







def 			__a			(     _UpperCamelCase:					ConvaD      ,		_UpperCamelCase:					torch.LongTensor      ,		_UpperCamelCase:					int = 1					)		->				ConvaD:

     """simple docstring"""




     _snake_case							=   index.to(layer.weight.device					)
     _snake_case							=   layer.weight.index_select(_UpperCamelCase      ,		_UpperCamelCase					).clone().detach()
     if dim == 0:
          _snake_case							=   layer.bias.clone().detach()
     else:
          _snake_case							=   layer.bias[index].clone().detach()
     _snake_case							=   list(layer.weight.size()					)
     _snake_case							=   len(_UpperCamelCase					)
     _snake_case							=   ConvaD(new_size[1]      ,		new_size[0]					).to(layer.weight.device					)
     _snake_case							=   False
     new_layer.weight.copy_(W.contiguous()					)
     _snake_case							=   True
     _snake_case							=   False
     new_layer.bias.copy_(b.contiguous()					)
     _snake_case							=   True
     return new_layer





def 			__a			(     _UpperCamelCase:					Union[nn.Linear, ConvaD]      ,		_UpperCamelCase:					torch.LongTensor      ,		_UpperCamelCase:					Optional[int] = None					)		->				Union[nn.Linear, ConvaD]:

     """simple docstring"""




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





def 			__a			(     _UpperCamelCase:					Callable[..., torch.Tensor]      ,		_UpperCamelCase:					int      ,		_UpperCamelCase:					int      ,		*_UpperCamelCase:					int					)		->				torch.Tensor:

     """simple docstring"""




     assert len(_UpperCamelCase					) > 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
     _snake_case							=   len(inspect.signature(_UpperCamelCase					).parameters					)
     if num_args_in_forward_chunk_fn != len(_UpperCamelCase					):
          raise ValueError(
              F"""forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(_UpperCamelCase					)} input """
              "tensors are given"					)

     if chunk_size > 0:
          _snake_case							=   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}"""					)

          _snake_case							=   input_tensors[0].shape[chunk_dim] // chunk_size

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

     return forward_fn(*_UpperCamelCase					)





def 			__a			(     _UpperCamelCase:					List[int]      ,		_UpperCamelCase:					int      ,		_UpperCamelCase:					int      ,		_UpperCamelCase:					Set[int]					)		->				Tuple[Set[int], torch.LongTensor]:

     """simple docstring"""




     _snake_case							=   torch.ones(_UpperCamelCase      ,		_UpperCamelCase					)
     _snake_case							=   set(_UpperCamelCase					) - 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
          _snake_case							=   head - sum(1 if h < head else 0 for h in already_pruned_heads					)
          _snake_case							=   0
     _snake_case							=   mask.view(-1					).contiguous().eq(1					)
     _snake_case							=   torch.arange(len(_UpperCamelCase					)					)[mask].long()
     return heads, index





def 			__a			(     *_UpperCamelCase:					Union[torch.Tensor, List[torch.Tensor]]      ,		_UpperCamelCase:					Optional[str] = None					)		->				Tuple[torch.Tensor, ...]:

     """simple docstring"""




     return torch.meshgrid(*_UpperCamelCase      ,		indexing=_UpperCamelCase					)





def 			__a			(     _UpperCamelCase:					torch.Tensor					)		->				Tuple[torch.device, int, int]:

     """simple docstring"""




     return tensor.device, storage_ptr(_UpperCamelCase					), storage_size(_UpperCamelCase					)








'''simple docstring'''





import argparse
import os


UpperCamelCase_		:    Union[str, Any]									=		{
    '''pytorch''': {
        '''1.13''': {
            '''torch''': '''1.13.1''',
            '''torchvision''': '''0.14.1''',
            '''torchaudio''': '''0.13.1''',
            '''python''': 3.9,
            '''cuda''': '''cu116''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu116'''
            ),
            '''base_image''': '''nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04''',
        },
        '''1.12''': {
            '''torch''': '''1.12.1''',
            '''torchvision''': '''0.13.1''',
            '''torchaudio''': '''0.12.1''',
            '''python''': 3.9,
            '''cuda''': '''cu113''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu113'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''1.11''': {
            '''torch''': '''1.11.0''',
            '''torchvision''': '''0.12.0''',
            '''torchaudio''': '''0.11.0''',
            '''python''': 3.9,
            '''cuda''': '''cu113''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu113'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''1.10''': {
            '''torch''': '''1.10.2''',
            '''torchvision''': '''0.11.3''',
            '''torchaudio''': '''0.10.2''',
            '''python''': 3.9,
            '''cuda''': '''cu113''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.10.2 torchvision==0.11.3 torchaudio==0.10.2'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu113'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        # torchaudio < 0.10 has no CUDA-enabled binary distributions
        '''1.9''': {
            '''torch''': '''1.9.1''',
            '''torchvision''': '''0.10.1''',
            '''torchaudio''': '''0.9.1''',
            '''python''': 3.9,
            '''cuda''': '''cu111''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.9.1 torchvision==0.10.1 torchaudio==0.9.1'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu111'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
    },
    '''tensorflow''': {
        '''2.11''': {
            '''tensorflow''': '''2.11.1''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.11.1''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.10''': {
            '''tensorflow''': '''2.10.1''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.10.1''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.9''': {
            '''tensorflow''': '''2.9.3''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.9.3''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.8''': {
            '''tensorflow''': '''2.8.2''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.8.2''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.7''': {
            '''tensorflow''': '''2.7.3''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.7.3''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.6''': {
            '''tensorflow''': '''2.6.5''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.6.5''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.5''': {
            '''tensorflow''': '''2.5.3''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.5.3''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
    },
}


if __name__ == "__main__":
   UpperCamelCase_		:    Optional[Any]									=		argparse.ArgumentParser('''Choose the framework and version to install''')
   parser.add_argument(
       '''--framework''', help='''The framework to install. Should be `torch` or `tensorflow`''', type=str, required=True
   )
   parser.add_argument('''--version''', help='''The version of the framework to install.''', type=str, required=True)
   UpperCamelCase_		:    Tuple									=		parser.parse_args()

   UpperCamelCase_		:    Optional[int]									=		past_versions_testing[args.framework][args.version]

   os.system(F'echo "export INSTALL_CMD=\'{info["install"]}\'" >> ~/.profile')
   print(F'echo "export INSTALL_CMD=\'{info["install"]}\'" >> ~/.profile')

   UpperCamelCase_		:    List[Any]									=		''''''
   if args.framework == "pytorch":
      UpperCamelCase_		:    str									=		info['''cuda''']
   os.system(F'echo \"export CUDA=\'{cuda}\'\" >> ~/.profile')
   print(F'echo \"export CUDA=\'{cuda}\'\" >> ~/.profile')


