

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


_a	 =     [nn.LayerNorm]

_a	 =     logging.get_logger(__name__)

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

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







def _a				( SCREAMING_SNAKE_CASE	:    List[Any]					,      SCREAMING_SNAKE_CASE	:    int					,      SCREAMING_SNAKE_CASE	:    List[Any]					,      SCREAMING_SNAKE_CASE	:    int					,      self	:    str	)					->					Optional[Any]:

	"""simple docstring"""






	from torch import _softmax_backward_data

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







def _a				( SCREAMING_SNAKE_CASE	:    nn.Linear					,      SCREAMING_SNAKE_CASE	:    torch.LongTensor					,      SCREAMING_SNAKE_CASE	:    int = 0	)					->					nn.Linear:

	"""simple docstring"""






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






class 							A_				(	nn.Module      ):


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


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







def _a				( SCREAMING_SNAKE_CASE	:    ConvaD					,      SCREAMING_SNAKE_CASE	:    torch.LongTensor					,      SCREAMING_SNAKE_CASE	:    int = 1	)					->					ConvaD:

	"""simple docstring"""






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







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

	"""simple docstring"""






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







def _a				( SCREAMING_SNAKE_CASE	:    Callable[..., torch.Tensor]					,      SCREAMING_SNAKE_CASE	:    int					,      SCREAMING_SNAKE_CASE	:    int					,      *SCREAMING_SNAKE_CASE	:    Dict	)					->					torch.Tensor:

	"""simple docstring"""






	assert len(SCREAMING_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
	__lowerCAmelCase:     Any     			=   len(inspect.signature(SCREAMING_SNAKE_CASE	).parameters	)
	if num_args_in_forward_chunk_fn != len(SCREAMING_SNAKE_CASE	):
		raise ValueError(
		    f'''forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(SCREAMING_SNAKE_CASE	)} input '''
		    'tensors are given'	)

	if chunk_size > 0:
		__lowerCAmelCase:     Optional[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}'''	)

		__lowerCAmelCase:     List[Any]     			=   input_tensors[0].shape[chunk_dim] // chunk_size

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

	return forward_fn(*SCREAMING_SNAKE_CASE	)







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

	"""simple docstring"""






	__lowerCAmelCase:     Optional[int]     			=   torch.ones(SCREAMING_SNAKE_CASE					,      SCREAMING_SNAKE_CASE	)
	__lowerCAmelCase:     Optional[int]     			=   set(SCREAMING_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
		__lowerCAmelCase:     List[str]     			=   head - sum(1 if h < head else 0 for h in already_pruned_heads	)
		__lowerCAmelCase:     Dict     			=   0
	__lowerCAmelCase:     Any     			=   mask.view(-1	).contiguous().eq(1	)
	__lowerCAmelCase:     torch.LongTensor     			=   torch.arange(len(SCREAMING_SNAKE_CASE	)	)[mask].long()
	return heads, index







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

	"""simple docstring"""






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







def _a				( SCREAMING_SNAKE_CASE	:    torch.Tensor	)					->					Tuple[torch.device, int, int]:

	"""simple docstring"""






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








# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
    is_vision_available,
)


_a	 =     {
    '''configuration_tvlt''': ['''TVLT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''TvltConfig'''],
    '''feature_extraction_tvlt''': ['''TvltFeatureExtractor'''],
    '''processing_tvlt''': ['''TvltProcessor'''],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _a	 =     [
      '''TVLT_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''TvltModel''',
      '''TvltForPreTraining''',
      '''TvltForAudioVisualClassification''',
      '''TvltPreTrainedModel''',
  ]

try:
  if not is_vision_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _a	 =     ['''TvltImageProcessor''']


if TYPE_CHECKING:
  from .configuration_tvlt import TVLT_PRETRAINED_CONFIG_ARCHIVE_MAP, TvltConfig
  from .processing_tvlt import TvltProcessor
  from .feature_extraction_tvlt import TvltFeatureExtractor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tvlt import (
        TVLT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TvltForAudioVisualClassification,
        TvltForPreTraining,
        TvltModel,
        TvltPreTrainedModel,
    )

  try:
    if not is_vision_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .image_processing_tvlt import TvltImageProcessor


else:
  import sys

  _a	 =     _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





