






import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    MODEL_FOR_PRETRAINING_MAPPING,
							    FunnelBaseModel,
							    FunnelForMaskedLM,
							    FunnelForMultipleChoice,
							    FunnelForPreTraining,
							    FunnelForQuestionAnswering,
							    FunnelForSequenceClassification,
							    FunnelForTokenClassification,
							    FunnelModel,
							)


class 					__SCREAMING_SNAKE_CASE:
						def __init__(       self: Optional[Any]    ,		UpperCamelCase: Tuple    ,		UpperCamelCase: Any=13    ,		UpperCamelCase: Union[str, Any]=7    ,		UpperCamelCase: int=True    ,		UpperCamelCase: Union[str, Any]=True    ,		UpperCamelCase: int=True    ,		UpperCamelCase: List[str]=True    ,		UpperCamelCase: Union[str, Any]=99    ,		UpperCamelCase: List[Any]=[1, 1, 2]    ,		UpperCamelCase: Union[str, Any]=1    ,		UpperCamelCase: Dict=32    ,		UpperCamelCase: str=4    ,		UpperCamelCase: Tuple=8    ,		UpperCamelCase: Any=37    ,		UpperCamelCase: List[Any]="gelu_new"    ,		UpperCamelCase: Optional[int]=0.1    ,		UpperCamelCase: Optional[Any]=0.1    ,		UpperCamelCase: Optional[Any]=0.0    ,		UpperCamelCase: Dict=5_12    ,		UpperCamelCase: Optional[int]=3    ,		UpperCamelCase: str=0.02    ,		UpperCamelCase: List[Any]=3    ,		UpperCamelCase: List[str]=4    ,		UpperCamelCase: List[Any]=None    ,		UpperCamelCase: str=False    ,		)						->						Optional[Any]:
										snake_case__													=       parent
										snake_case__													=       batch_size
										snake_case__													=       seq_length
										snake_case__													=       is_training
										snake_case__													=       use_input_mask
										snake_case__													=       use_token_type_ids
										snake_case__													=       use_labels
										snake_case__													=       vocab_size
										snake_case__													=       block_sizes
										snake_case__													=       num_decoder_layers
										snake_case__													=       d_model
										snake_case__													=       n_head
										snake_case__													=       d_head
										snake_case__													=       d_inner
										snake_case__													=       hidden_act
										snake_case__													=       hidden_dropout
										snake_case__													=       attention_dropout
										snake_case__													=       activation_dropout
										snake_case__													=       max_position_embeddings
										snake_case__													=       type_vocab_size
										snake_case__													=       2
										snake_case__													=       num_labels
										snake_case__													=       num_choices
										snake_case__													=       scope
										snake_case__													=       initializer_std

										# Used in the tests to check the size of the first attention layer
										snake_case__													=       n_head
										# Used in the tests to check the size of the first hidden state
										snake_case__													=       self.d_model
										# Used in the tests to check the number of output hidden states/attentions
										snake_case__													=       sum(self.block_sizes  ) + (0 if base else self.num_decoder_layers)
										# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
										# the last hidden state of the first block (which is the first hidden state of the decoder).
										if not base:
														snake_case__													=       self.num_hidden_layers + 2
						def lowerCAmelCase_	(       self: Optional[int]  )						->						Dict:
										snake_case__													=       ids_tensor([self.batch_size, self.seq_length]    ,		self.vocab_size  )

										snake_case__													=       None
										if self.use_input_mask:
														snake_case__													=       ids_tensor([self.batch_size, self.seq_length]    ,		vocab_size=2  )

										snake_case__													=       None
										if self.use_token_type_ids:
														snake_case__													=       ids_tensor([self.batch_size, self.seq_length]    ,		self.type_vocab_size  )

										snake_case__													=       None
										snake_case__													=       None
										snake_case__													=       None
										if self.use_labels:
														snake_case__													=       ids_tensor([self.batch_size]    ,		self.type_sequence_label_size  )
														snake_case__													=       ids_tensor([self.batch_size, self.seq_length]    ,		self.num_labels  )
														snake_case__													=       ids_tensor([self.batch_size]    ,		self.num_choices  )
														snake_case__													=       ids_tensor([self.batch_size, self.seq_length]    ,		1  )

										snake_case__													=       self.get_config()

										return (
										    config,
										    input_ids,
										    token_type_ids,
										    input_mask,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										    fake_token_labels,
										)
						def lowerCAmelCase_	(       self: Dict  )						->						Any:
										return FunnelConfig(
										    vocab_size=self.vocab_size    ,		block_sizes=self.block_sizes    ,		num_decoder_layers=self.num_decoder_layers    ,		d_model=self.d_model    ,		n_head=self.n_head    ,		d_head=self.d_head    ,		d_inner=self.d_inner    ,		hidden_act=self.hidden_act    ,		hidden_dropout=self.hidden_dropout    ,		attention_dropout=self.attention_dropout    ,		activation_dropout=self.activation_dropout    ,		max_position_embeddings=self.max_position_embeddings    ,		type_vocab_size=self.type_vocab_size    ,		initializer_std=self.initializer_std    ,		)
						def lowerCAmelCase_	(       self: List[str]    ,		UpperCamelCase: Dict    ,		UpperCamelCase: Tuple    ,		UpperCamelCase: List[str]    ,		UpperCamelCase: List[str]    ,		UpperCamelCase: int    ,		UpperCamelCase: Tuple    ,		UpperCamelCase: Optional[int]    ,		UpperCamelCase: Dict    ,		)						->						Any:
										snake_case__													=       FunnelModel(config=UpperCamelCase  )
										model.to(UpperCamelCase  )
										model.eval()
										snake_case__													=       model(UpperCamelCase    ,		attention_mask=UpperCamelCase    ,		token_type_ids=UpperCamelCase  )
										snake_case__													=       model(UpperCamelCase    ,		token_type_ids=UpperCamelCase  )
										snake_case__													=       model(UpperCamelCase  )
										self.parent.assertEqual(result.last_hidden_state.shape    ,		(self.batch_size, self.seq_length, self.d_model)  )

										snake_case__													=       False
										snake_case__													=       model(UpperCamelCase  )
										self.parent.assertEqual(result.last_hidden_state.shape    ,		(self.batch_size, self.seq_length, self.d_model)  )

										snake_case__													=       False
										snake_case__													=       model(UpperCamelCase  )
										self.parent.assertEqual(result.last_hidden_state.shape    ,		(self.batch_size, self.seq_length, self.d_model)  )
						def lowerCAmelCase_	(       self: List[Any]    ,		UpperCamelCase: Any    ,		UpperCamelCase: List[str]    ,		UpperCamelCase: Any    ,		UpperCamelCase: int    ,		UpperCamelCase: List[Any]    ,		UpperCamelCase: str    ,		UpperCamelCase: Any    ,		UpperCamelCase: Dict    ,		)						->						Union[str, Any]:
										snake_case__													=       FunnelBaseModel(config=UpperCamelCase  )
										model.to(UpperCamelCase  )
										model.eval()
										snake_case__													=       model(UpperCamelCase    ,		attention_mask=UpperCamelCase    ,		token_type_ids=UpperCamelCase  )
										snake_case__													=       model(UpperCamelCase    ,		token_type_ids=UpperCamelCase  )
										snake_case__													=       model(UpperCamelCase  )
										self.parent.assertEqual(result.last_hidden_state.shape    ,		(self.batch_size, 2, self.d_model)  )

										snake_case__													=       False
										snake_case__													=       model(UpperCamelCase  )
										self.parent.assertEqual(result.last_hidden_state.shape    ,		(self.batch_size, 3, self.d_model)  )

										snake_case__													=       False
										snake_case__													=       model(UpperCamelCase  )
										self.parent.assertEqual(result.last_hidden_state.shape    ,		(self.batch_size, 2, self.d_model)  )
						def lowerCAmelCase_	(       self: Optional[int]    ,		UpperCamelCase: Union[str, Any]    ,		UpperCamelCase: Dict    ,		UpperCamelCase: Dict    ,		UpperCamelCase: int    ,		UpperCamelCase: Tuple    ,		UpperCamelCase: List[Any]    ,		UpperCamelCase: Optional[int]    ,		UpperCamelCase: str    ,		)						->						Any:
										snake_case__													=       self.num_labels
										snake_case__													=       FunnelForPreTraining(config=UpperCamelCase  )
										model.to(UpperCamelCase  )
										model.eval()
										snake_case__													=       model(UpperCamelCase    ,		attention_mask=UpperCamelCase    ,		token_type_ids=UpperCamelCase    ,		labels=UpperCamelCase  )
										self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.seq_length)  )
						def lowerCAmelCase_	(       self: Any    ,		UpperCamelCase: Dict    ,		UpperCamelCase: Dict    ,		UpperCamelCase: str    ,		UpperCamelCase: Optional[int]    ,		UpperCamelCase: Optional[int]    ,		UpperCamelCase: Tuple    ,		UpperCamelCase: Optional[Any]    ,		UpperCamelCase: List[Any]    ,		)						->						Optional[int]:
										snake_case__													=       FunnelForMaskedLM(config=UpperCamelCase  )
										model.to(UpperCamelCase  )
										model.eval()
										snake_case__													=       model(UpperCamelCase    ,		attention_mask=UpperCamelCase    ,		token_type_ids=UpperCamelCase    ,		labels=UpperCamelCase  )
										self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.seq_length, self.vocab_size)  )
						def lowerCAmelCase_	(       self: Tuple    ,		UpperCamelCase: Any    ,		UpperCamelCase: str    ,		UpperCamelCase: List[Any]    ,		UpperCamelCase: str    ,		UpperCamelCase: Dict    ,		UpperCamelCase: Optional[int]    ,		UpperCamelCase: Any    ,		UpperCamelCase: Dict    ,		)						->						List[Any]:
										snake_case__													=       self.num_labels
										snake_case__													=       FunnelForSequenceClassification(UpperCamelCase  )
										model.to(UpperCamelCase  )
										model.eval()
										snake_case__													=       model(UpperCamelCase    ,		attention_mask=UpperCamelCase    ,		token_type_ids=UpperCamelCase    ,		labels=UpperCamelCase  )
										self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.num_labels)  )
						def lowerCAmelCase_	(       self: Dict    ,		UpperCamelCase: int    ,		UpperCamelCase: Optional[int]    ,		UpperCamelCase: Union[str, Any]    ,		UpperCamelCase: Union[str, Any]    ,		UpperCamelCase: Any    ,		UpperCamelCase: Optional[int]    ,		UpperCamelCase: List[Any]    ,		UpperCamelCase: Optional[Any]    ,		)						->						List[Any]:
										snake_case__													=       self.num_choices
										snake_case__													=       FunnelForMultipleChoice(config=UpperCamelCase  )
										model.to(UpperCamelCase  )
										model.eval()
										snake_case__													=       input_ids.unsqueeze(1  ).expand(-1    ,		self.num_choices    ,		-1  ).contiguous()
										snake_case__													=       token_type_ids.unsqueeze(1  ).expand(-1    ,		self.num_choices    ,		-1  ).contiguous()
										snake_case__													=       input_mask.unsqueeze(1  ).expand(-1    ,		self.num_choices    ,		-1  ).contiguous()
										snake_case__													=       model(
										    UpperCamelCase    ,		attention_mask=UpperCamelCase    ,		token_type_ids=UpperCamelCase    ,		labels=UpperCamelCase    ,		)
										self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.num_choices)  )
						def lowerCAmelCase_	(       self: Union[str, Any]    ,		UpperCamelCase: Any    ,		UpperCamelCase: List[Any]    ,		UpperCamelCase: Any    ,		UpperCamelCase: List[Any]    ,		UpperCamelCase: Dict    ,		UpperCamelCase: Dict    ,		UpperCamelCase: Dict    ,		UpperCamelCase: List[str]    ,		)						->						Tuple:
										snake_case__													=       self.num_labels
										snake_case__													=       FunnelForTokenClassification(config=UpperCamelCase  )
										model.to(UpperCamelCase  )
										model.eval()
										snake_case__													=       model(UpperCamelCase    ,		attention_mask=UpperCamelCase    ,		token_type_ids=UpperCamelCase    ,		labels=UpperCamelCase  )
										self.parent.assertEqual(result.logits.shape    ,		(self.batch_size, self.seq_length, self.num_labels)  )
						def lowerCAmelCase_	(       self: Any    ,		UpperCamelCase: List[Any]    ,		UpperCamelCase: Dict    ,		UpperCamelCase: List[str]    ,		UpperCamelCase: Any    ,		UpperCamelCase: Optional[Any]    ,		UpperCamelCase: Union[str, Any]    ,		UpperCamelCase: List[Any]    ,		UpperCamelCase: Any    ,		)						->						str:
										snake_case__													=       FunnelForQuestionAnswering(config=UpperCamelCase  )
										model.to(UpperCamelCase  )
										model.eval()
										snake_case__													=       model(
										    UpperCamelCase    ,		attention_mask=UpperCamelCase    ,		token_type_ids=UpperCamelCase    ,		start_positions=UpperCamelCase    ,		end_positions=UpperCamelCase    ,		)
										self.parent.assertEqual(result.start_logits.shape    ,		(self.batch_size, self.seq_length)  )
										self.parent.assertEqual(result.end_logits.shape    ,		(self.batch_size, self.seq_length)  )


						def lowerCAmelCase_	(       self: Tuple  )						->						str:
										snake_case__													=       self.prepare_config_and_inputs()
										(
										    snake_case__						
										)													=       config_and_inputs
										snake_case__													=       {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
										return config, inputs_dict


@require_torch
class 					__SCREAMING_SNAKE_CASE(		UpperCAmelCase__      ,					UpperCAmelCase__      ,					unittest.TestCase		):
						_UpperCAmelCase	      =    False
						_UpperCAmelCase	      =    False
						_UpperCAmelCase	      =    (
						    (
						        FunnelModel,
						        FunnelForMaskedLM,
						        FunnelForPreTraining,
						        FunnelForQuestionAnswering,
						        FunnelForTokenClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						_UpperCAmelCase	      =    (
						    {
						        'feature-extraction': (FunnelBaseModel, FunnelModel),
						        'fill-mask': FunnelForMaskedLM,
						        'question-answering': FunnelForQuestionAnswering,
						        'text-classification': FunnelForSequenceClassification,
						        'token-classification': FunnelForTokenClassification,
						        'zero-shot': FunnelForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						def lowerCAmelCase_	(       self: Optional[Any]    ,		UpperCamelCase: str    ,		UpperCamelCase: Optional[int]    ,		UpperCamelCase: Union[str, Any]=False  )						->						Optional[Any]:
										snake_case__													=       super()._prepare_for_class(UpperCamelCase    ,		UpperCamelCase    ,		return_labels=UpperCamelCase  )

										if return_labels:
														if model_class in get_values(UpperCamelCase  ):
																		snake_case__													=       torch.zeros(
																		    (self.model_tester.batch_size, self.model_tester.seq_length)    ,		dtype=torch.long    ,		device=UpperCamelCase  )
										return inputs_dict
						def lowerCAmelCase_	(       self: int  )						->						int:
										snake_case__													=       FunnelModelTester(self  )
										snake_case__													=       ConfigTester(self    ,		config_class=UpperCamelCase  )
						def lowerCAmelCase_	(       self: Optional[int]  )						->						Optional[int]:
										self.config_tester.run_common_tests()
						def lowerCAmelCase_	(       self: List[Any]  )						->						str:
										snake_case__													=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*UpperCamelCase  )
						def lowerCAmelCase_	(       self: List[str]  )						->						Any:
										snake_case__													=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_pretraining(*UpperCamelCase  )
						def lowerCAmelCase_	(       self: str  )						->						Optional[int]:
										snake_case__													=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase  )
						def lowerCAmelCase_	(       self: List[Any]  )						->						Dict:
										snake_case__													=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*UpperCamelCase  )
						def lowerCAmelCase_	(       self: Optional[int]  )						->						List[Any]:
										snake_case__													=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*UpperCamelCase  )


						def lowerCAmelCase_	(       self: List[str]    ,		UpperCamelCase: List[str]  )						->						Tuple:




										if hasattr(UpperCamelCase    ,		'weight'  ) and module.weight is not None:
														module.weight.data.fill_(3  )
										if hasattr(UpperCamelCase    ,		'bias'  ) and module.bias is not None:
														module.bias.data.fill_(3  )

										for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
														if hasattr(UpperCamelCase    ,		UpperCamelCase  ) and getattr(UpperCamelCase    ,		UpperCamelCase  ) is not None:
																		snake_case__													=       getattr(UpperCamelCase    ,		UpperCamelCase  )
																		weight.data.fill_(3  )


@require_torch
class 					__SCREAMING_SNAKE_CASE(		UpperCAmelCase__      ,					unittest.TestCase		):
						_UpperCAmelCase	      =    False
						_UpperCAmelCase	      =    False
						_UpperCAmelCase	      =    (
						    (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
						)
						def lowerCAmelCase_	(       self: Optional[Any]  )						->						Union[str, Any]:
										snake_case__													=       FunnelModelTester(self    ,		base=UpperCamelCase  )
										snake_case__													=       ConfigTester(self    ,		config_class=UpperCamelCase  )
						def lowerCAmelCase_	(       self: Optional[int]  )						->						List[Any]:
										self.config_tester.run_common_tests()
						def lowerCAmelCase_	(       self: int  )						->						Any:
										snake_case__													=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_base_model(*UpperCamelCase  )
						def lowerCAmelCase_	(       self: Dict  )						->						Union[str, Any]:
										snake_case__													=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase  )
						def lowerCAmelCase_	(       self: List[Any]  )						->						Optional[int]:
										snake_case__													=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase  )
						def lowerCAmelCase_	(       self: Tuple  )						->						str:
										snake_case__													=       self.model_tester.prepare_config_and_inputs_for_common()
										snake_case__													=       True

										for model_class in self.all_model_classes:
														if model_class.__name__ == "FunnelBaseModel":
																		continue
														snake_case__													=       model_class(UpperCamelCase  )
														model.to(UpperCamelCase  )
														model.train()
														snake_case__													=       self._prepare_for_class(UpperCamelCase    ,		UpperCamelCase    ,		return_labels=UpperCamelCase  )
														snake_case__													=       model(**UpperCamelCase  ).loss
														loss.backward()


						def lowerCAmelCase_	(       self: Optional[int]    ,		UpperCamelCase: str  )						->						Union[str, Any]:




										if hasattr(UpperCamelCase    ,		'weight'  ) and module.weight is not None:
														module.weight.data.fill_(3  )
										if hasattr(UpperCamelCase    ,		'bias'  ) and module.bias is not None:
														module.bias.data.fill_(3  )

										for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
														if hasattr(UpperCamelCase    ,		UpperCamelCase  ) and getattr(UpperCamelCase    ,		UpperCamelCase  ) is not None:
																		snake_case__													=       getattr(UpperCamelCase    ,		UpperCamelCase  )
																		weight.data.fill_(3  )


@require_torch
@require_sentencepiece
@require_tokenizers
class 					__SCREAMING_SNAKE_CASE(		unittest.TestCase		):
						def lowerCAmelCase_	(       self: Any  )						->						int:
										snake_case__													=       13
										snake_case__													=       7
										snake_case__													=       torch.arange(0    ,		batch_size * sequence_length  ).long().reshape(UpperCamelCase    ,		UpperCamelCase  )
										snake_case__													=       [0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
										snake_case__													=       torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]  )

										snake_case__													=       FunnelModel.from_pretrained('sgugger/funnel-random-tiny'  )
										snake_case__													=       model(UpperCamelCase    ,		token_type_ids=UpperCamelCase  )[0].abs()

										snake_case__													=       torch.tensor(23_44.83_52  )
										snake_case__													=       torch.tensor(0.8_052  )
										self.assertTrue(torch.allclose(output.sum()    ,		UpperCamelCase    ,		atol=1e-4  )  )
										self.assertTrue(torch.allclose(output.mean()    ,		UpperCamelCase    ,		atol=1e-4  )  )

										snake_case__													=       torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]  )
										snake_case__													=       model(UpperCamelCase    ,		attention_mask=UpperCamelCase    ,		token_type_ids=UpperCamelCase  )[0].abs()

										snake_case__													=       torch.tensor(23_43.84_25  )
										snake_case__													=       torch.tensor(0.8_049  )
										self.assertTrue(torch.allclose(output.sum()    ,		UpperCamelCase    ,		atol=1e-4  )  )
										self.assertTrue(torch.allclose(output.mean()    ,		UpperCamelCase    ,		atol=1e-4  )  )


						@slow
						def lowerCAmelCase_	(       self: List[Any]  )						->						Any:
										snake_case__													=       FunnelTokenizer.from_pretrained('huggingface/funnel-small'  )
										snake_case__													=       FunnelModel.from_pretrained('huggingface/funnel-small'  )
										snake_case__													=       tokenizer('Hello! I am the Funnel Transformer model.'    ,		return_tensors='pt'  )
										snake_case__													=       model(**UpperCamelCase  )[0]

										snake_case__													=       torch.tensor(2_35.72_46  )
										snake_case__													=       torch.tensor(0.0_256  )
										self.assertTrue(torch.allclose(output.sum()    ,		UpperCamelCase    ,		atol=1e-4  )  )
										self.assertTrue(torch.allclose(output.mean()    ,		UpperCamelCase    ,		atol=1e-4  )  )












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   :						int							=  [nn.LayerNorm]

__UpperCamelCase   :						List[str]							=  logging.get_logger(__name__)

__UpperCamelCase   :						str							=  version.parse(version.parse(torch.__version__).base_version)

__UpperCamelCase   :						List[str]							=  parsed_torch_version_base >= version.parse("""2.0""")
__UpperCamelCase   :						Dict							=  parsed_torch_version_base >= version.parse("""1.12""")
__UpperCamelCase   :						Tuple							=  parsed_torch_version_base >= version.parse("""1.11""")
__UpperCamelCase   :						Union[str, Any]							=  parsed_torch_version_base < version.parse("""1.11""")
def        a_	( _A ,				_A ,				_A ,				_A ,				self					)     ->     Any:







				"""simple docstring"""
				from torch import _softmax_backward_data

				if is_torch_less_than_1_11:
								return _softmax_backward_data(_A ,				_A ,				parent.dim ,				self					)
				else:
								return _softmax_backward_data(_A ,				_A ,				parent.dim ,				self.dtype					)
def        a_	( _A ,				_A ,				_A = 0					)     ->     nn.Linear:







				"""simple docstring"""
				snake_case__													=       index.to(layer.weight.device					)
				snake_case__													=       layer.weight.index_select(_A ,				_A					).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(_A					)
				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 					__SCREAMING_SNAKE_CASE(		nn.Module		):
							def __init__(       self: Union[str, Any]    ,		UpperCamelCase: Dict    ,		UpperCamelCase: List[Any]  )						->						Any:
											super().__init__()
											snake_case__													=       nf
											snake_case__													=       nn.Parameter(torch.empty(UpperCamelCase    ,		UpperCamelCase  )  )
											snake_case__													=       nn.Parameter(torch.zeros(UpperCamelCase  )  )
											nn.init.normal_(self.weight    ,		std=0.02  )


							def lowerCAmelCase_	(       self: Optional[Any]    ,		UpperCamelCase: int  )						->						List[Any]:
											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(UpperCamelCase  )
											return x


def        a_	( _A ,				_A ,				_A = 1					)     ->     ConvaD:







				"""simple docstring"""
				snake_case__													=       index.to(layer.weight.device					)
				snake_case__													=       layer.weight.index_select(_A ,				_A					).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(_A					)
				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_	( _A ,				_A ,				_A = None					)     ->     Union[nn.Linear, ConvaD]:







				"""simple docstring"""
				if isinstance(_A ,				nn.Linear					):
								return prune_linear_layer(_A ,				_A ,				dim=0 if dim is None else dim					)
				elif isinstance(_A ,				_A					):
								return prune_convad_layer(_A ,				_A ,				dim=1 if dim is None else dim					)
				else:
								raise ValueError(f'''Can\'t prune layer of class {layer.__class__}'''					)
def        a_	( _A ,				_A ,				_A ,				*_A					)     ->     torch.Tensor:







				"""simple docstring"""
				assert len(_A					) > 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(_A					).parameters					)
				if num_args_in_forward_chunk_fn != len(_A					):
								raise ValueError(
								    f'''forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(_A					)} 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(_A ,				dim=_A					) for input_tensor in input_tensors					)
								# apply forward fn to every tuple
								snake_case__													=       tuple(forward_fn(*_A					) for input_tensors_chunk in zip(*_A					)					)
								# concatenate output at same dimension
								return torch.cat(_A ,				dim=_A					)

				return forward_fn(*_A					)
def        a_	( _A ,				_A ,				_A ,				_A					)     ->     Tuple[Set[int], torch.LongTensor]:







				"""simple docstring"""
				snake_case__													=       torch.ones(_A ,				_A					)
				snake_case__													=       set(_A					) - 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(_A					)					)[mask].long()
				return heads, index
def        a_	( *_A ,				_A = None					)     ->     Tuple[torch.Tensor, ...]:







				"""simple docstring"""
				return torch.meshgrid(*_A ,				indexing=_A					)


def        a_	( _A					)     ->     Tuple[torch.device, int, int]:







				"""simple docstring"""
				return tensor.device, storage_ptr(_A					), storage_size(_A					)




