




import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
				import torch

				from transformers import (
				    MODEL_FOR_PRETRAINING_MAPPING,
				    ElectraForCausalLM,
				    ElectraForMaskedLM,
				    ElectraForMultipleChoice,
				    ElectraForPreTraining,
				    ElectraForQuestionAnswering,
				    ElectraForSequenceClassification,
				    ElectraForTokenClassification,
				    ElectraModel,
				)
				from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST




class     A					:







							def __init__(    self   :	List[Any]   ,     __UpperCAmelCase   :	Optional[Any]   ,     __UpperCAmelCase   :	str=13   ,     __UpperCAmelCase   :	Union[str, Any]=7   ,     __UpperCAmelCase   :	Tuple=True   ,     __UpperCAmelCase   :	Union[str, Any]=True   ,     __UpperCAmelCase   :	Any=True   ,     __UpperCAmelCase   :	str=True   ,     __UpperCAmelCase   :	str=99   ,     __UpperCAmelCase   :	int=32   ,     __UpperCAmelCase   :	Optional[int]=5   ,     __UpperCAmelCase   :	List[str]=4   ,     __UpperCAmelCase   :	List[Any]=37   ,     __UpperCAmelCase   :	Tuple="gelu"   ,     __UpperCAmelCase   :	Optional[Any]=0.1   ,     __UpperCAmelCase   :	Union[str, Any]=0.1   ,     __UpperCAmelCase   :	List[Any]=512   ,     __UpperCAmelCase   :	Dict=16   ,     __UpperCAmelCase   :	str=2   ,     __UpperCAmelCase   :	List[Any]=0.02   ,     __UpperCAmelCase   :	Tuple=3   ,     __UpperCAmelCase   :	Dict=4   ,     __UpperCAmelCase   :	List[Any]=None   ,     )		->					str:

										"""simple docstring"""



										UpperCamelCase_     				= parent
										UpperCamelCase_     				= batch_size
										UpperCamelCase_     				= seq_length
										UpperCamelCase_     				= is_training
										UpperCamelCase_     				= use_input_mask
										UpperCamelCase_     				= use_token_type_ids
										UpperCamelCase_     				= use_labels
										UpperCamelCase_     				= vocab_size
										UpperCamelCase_     				= hidden_size
										UpperCamelCase_     				= num_hidden_layers
										UpperCamelCase_     				= num_attention_heads
										UpperCamelCase_     				= intermediate_size
										UpperCamelCase_     				= hidden_act
										UpperCamelCase_     				= hidden_dropout_prob
										UpperCamelCase_     				= attention_probs_dropout_prob
										UpperCamelCase_     				= max_position_embeddings
										UpperCamelCase_     				= type_vocab_size
										UpperCamelCase_     				= type_sequence_label_size
										UpperCamelCase_     				= initializer_range
										UpperCamelCase_     				= num_labels
										UpperCamelCase_     				= num_choices
										UpperCamelCase_     				= scope







							def        lowercase__							(    self   :	Optional[Any]				)		->					List[str]:

										"""simple docstring"""



										UpperCamelCase_     				= ids_tensor([self.batch_size, self.seq_length]   ,     self.vocab_size				)

										UpperCamelCase_     				= None
										if self.use_input_mask:
													UpperCamelCase_     				= random_attention_mask([self.batch_size, self.seq_length]				)

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

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

										UpperCamelCase_     				= self.get_config()

										return (
										    config,
										    input_ids,
										    token_type_ids,
										    input_mask,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										    fake_token_labels,
										)







							def        lowercase__							(    self   :	str				)		->					Any:

										"""simple docstring"""



										return ElectraConfig(
										    vocab_size=self.vocab_size   ,     hidden_size=self.hidden_size   ,     num_hidden_layers=self.num_hidden_layers   ,     num_attention_heads=self.num_attention_heads   ,     intermediate_size=self.intermediate_size   ,     hidden_act=self.hidden_act   ,     hidden_dropout_prob=self.hidden_dropout_prob   ,     attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,     max_position_embeddings=self.max_position_embeddings   ,     type_vocab_size=self.type_vocab_size   ,     is_decoder=__UpperCAmelCase   ,     initializer_range=self.initializer_range   ,     )







							def        lowercase__							(    self   :	List[Any]				)		->					Union[str, Any]:

										"""simple docstring"""



										(
										    (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   
										)     				= self.prepare_config_and_inputs()

										UpperCamelCase_     				= True
										UpperCamelCase_     				= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
										UpperCamelCase_     				= ids_tensor([self.batch_size, self.seq_length]   ,     vocab_size=2				)

										return (
										    config,
										    input_ids,
										    token_type_ids,
										    input_mask,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										    encoder_hidden_states,
										    encoder_attention_mask,
										)







							def        lowercase__							(    self   :	Union[str, Any]   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Optional[Any]   ,     __UpperCAmelCase   :	Dict   ,     __UpperCAmelCase   :	Tuple   ,     __UpperCAmelCase   :	Union[str, Any]   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Optional[int]   ,     )		->					Tuple:

										"""simple docstring"""



										UpperCamelCase_     				= ElectraModel(config=__UpperCAmelCase				)
										model.to(__UpperCAmelCase				)
										model.eval()
										UpperCamelCase_     				= model(__UpperCAmelCase   ,     attention_mask=__UpperCAmelCase   ,     token_type_ids=__UpperCAmelCase				)
										UpperCamelCase_     				= model(__UpperCAmelCase   ,     token_type_ids=__UpperCAmelCase				)
										UpperCamelCase_     				= model(__UpperCAmelCase				)
										self.parent.assertEqual(result.last_hidden_state.shape   ,     (self.batch_size, self.seq_length, self.hidden_size)				)







							def        lowercase__							(    self   :	int   ,     __UpperCAmelCase   :	Optional[int]   ,     __UpperCAmelCase   :	Any   ,     __UpperCAmelCase   :	str   ,     __UpperCAmelCase   :	Optional[Any]   ,     __UpperCAmelCase   :	Dict   ,     __UpperCAmelCase   :	Any   ,     __UpperCAmelCase   :	Any   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	Tuple   ,     )		->					int:

										"""simple docstring"""



										UpperCamelCase_     				= True
										UpperCamelCase_     				= ElectraModel(__UpperCAmelCase				)
										model.to(__UpperCAmelCase				)
										model.eval()
										UpperCamelCase_     				= model(
										    __UpperCAmelCase   ,     attention_mask=__UpperCAmelCase   ,     token_type_ids=__UpperCAmelCase   ,     encoder_hidden_states=__UpperCAmelCase   ,     encoder_attention_mask=__UpperCAmelCase   ,     )
										UpperCamelCase_     				= model(
										    __UpperCAmelCase   ,     attention_mask=__UpperCAmelCase   ,     token_type_ids=__UpperCAmelCase   ,     encoder_hidden_states=__UpperCAmelCase   ,     )
										UpperCamelCase_     				= model(__UpperCAmelCase   ,     attention_mask=__UpperCAmelCase   ,     token_type_ids=__UpperCAmelCase				)
										self.parent.assertEqual(result.last_hidden_state.shape   ,     (self.batch_size, self.seq_length, self.hidden_size)				)







							def        lowercase__							(    self   :	List[str]   ,     __UpperCAmelCase   :	List[Any]   ,     __UpperCAmelCase   :	List[Any]   ,     __UpperCAmelCase   :	str   ,     __UpperCAmelCase   :	Optional[int]   ,     __UpperCAmelCase   :	List[Any]   ,     __UpperCAmelCase   :	Dict   ,     __UpperCAmelCase   :	Union[str, Any]   ,     __UpperCAmelCase   :	Tuple   ,     )		->					Optional[Any]:

										"""simple docstring"""



										UpperCamelCase_     				= ElectraForMaskedLM(config=__UpperCAmelCase				)
										model.to(__UpperCAmelCase				)
										model.eval()
										UpperCamelCase_     				= 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        lowercase__							(    self   :	Union[str, Any]   ,     __UpperCAmelCase   :	Optional[Any]   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	List[Any]   ,     __UpperCAmelCase   :	Any   ,     __UpperCAmelCase   :	str   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Dict   ,     __UpperCAmelCase   :	Tuple   ,     __UpperCAmelCase   :	List[str]   ,     )		->					Tuple:

										"""simple docstring"""



										UpperCamelCase_     				= ElectraForCausalLM(config=__UpperCAmelCase				)
										model.to(__UpperCAmelCase				)
										model.eval()
										UpperCamelCase_     				= 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        lowercase__							(    self   :	Union[str, Any]   ,     __UpperCAmelCase   :	str   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Optional[int]   ,     __UpperCAmelCase   :	Optional[Any]   ,     __UpperCAmelCase   :	Tuple   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Any   ,     )		->					Optional[int]:

										"""simple docstring"""



										UpperCamelCase_     				= self.num_labels
										UpperCamelCase_     				= ElectraForTokenClassification(config=__UpperCAmelCase				)
										model.to(__UpperCAmelCase				)
										model.eval()
										UpperCamelCase_     				= 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        lowercase__							(    self   :	List[str]   ,     __UpperCAmelCase   :	Union[str, Any]   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	List[Any]   ,     __UpperCAmelCase   :	Union[str, Any]   ,     __UpperCAmelCase   :	str   ,     __UpperCAmelCase   :	Tuple   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	str   ,     )		->					str:

										"""simple docstring"""



										UpperCamelCase_     				= self.num_labels
										UpperCamelCase_     				= ElectraForPreTraining(config=__UpperCAmelCase				)
										model.to(__UpperCAmelCase				)
										model.eval()
										UpperCamelCase_     				= model(__UpperCAmelCase   ,     attention_mask=__UpperCAmelCase   ,     token_type_ids=__UpperCAmelCase   ,     labels=__UpperCAmelCase				)
										self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.seq_length)				)







							def        lowercase__							(    self   :	Union[str, Any]   ,     __UpperCAmelCase   :	Tuple   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Dict   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	List[Any]   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Any   ,     )		->					Any:

										"""simple docstring"""



										UpperCamelCase_     				= self.num_labels
										UpperCamelCase_     				= ElectraForSequenceClassification(__UpperCAmelCase				)
										model.to(__UpperCAmelCase				)
										model.eval()
										UpperCamelCase_     				= model(__UpperCAmelCase   ,     attention_mask=__UpperCAmelCase   ,     token_type_ids=__UpperCAmelCase   ,     labels=__UpperCAmelCase				)
										self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.num_labels)				)







							def        lowercase__							(    self   :	Dict   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Tuple   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	List[Any]   ,     __UpperCAmelCase   :	Optional[Any]   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	List[Any]   ,     )		->					List[Any]:

										"""simple docstring"""



										UpperCamelCase_     				= ElectraForQuestionAnswering(config=__UpperCAmelCase				)
										model.to(__UpperCAmelCase				)
										model.eval()
										UpperCamelCase_     				= 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        lowercase__							(    self   :	Any   ,     __UpperCAmelCase   :	Dict   ,     __UpperCAmelCase   :	str   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	List[Any]   ,     __UpperCAmelCase   :	Optional[Any]   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	int   ,     __UpperCAmelCase   :	Dict   ,     )		->					Tuple:

										"""simple docstring"""



										UpperCamelCase_     				= self.num_choices
										UpperCamelCase_     				= ElectraForMultipleChoice(config=__UpperCAmelCase				)
										model.to(__UpperCAmelCase				)
										model.eval()
										UpperCamelCase_     				= input_ids.unsqueeze(1				).expand(-1   ,     self.num_choices   ,     -1				).contiguous()
										UpperCamelCase_     				= token_type_ids.unsqueeze(1				).expand(-1   ,     self.num_choices   ,     -1				).contiguous()
										UpperCamelCase_     				= input_mask.unsqueeze(1				).expand(-1   ,     self.num_choices   ,     -1				).contiguous()
										UpperCamelCase_     				= model(
										    __UpperCAmelCase   ,     attention_mask=__UpperCAmelCase   ,     token_type_ids=__UpperCAmelCase   ,     labels=__UpperCAmelCase   ,     )
										self.parent.assertEqual(result.logits.shape   ,     (self.batch_size, self.num_choices)				)







							def        lowercase__							(    self   :	List[str]				)		->					Optional[Any]:

										"""simple docstring"""



										UpperCamelCase_     				= self.prepare_config_and_inputs()
										(
										    (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   (
										    UpperCamelCase_
										)				,   
										)     				= config_and_inputs
										UpperCamelCase_     				= {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
										return config, inputs_dict







@require_torch
class     A					(				lowerCamelCase_  ,						lowerCamelCase_  ,						unittest.TestCase       ):
							_SCREAMING_SNAKE_CASE  :		Dict			  =       (
							    (
							        ElectraModel,
							        ElectraForPreTraining,
							        ElectraForMaskedLM,
							        ElectraForCausalLM,
							        ElectraForMultipleChoice,
							        ElectraForTokenClassification,
							        ElectraForSequenceClassification,
							        ElectraForQuestionAnswering,
							    )
							    if is_torch_available()
							    else ()
							)
							_SCREAMING_SNAKE_CASE  :		str			  =       (
							    {
							        '''feature-extraction''': ElectraModel,
							        '''fill-mask''': ElectraForMaskedLM,
							        '''question-answering''': ElectraForQuestionAnswering,
							        '''text-classification''': ElectraForSequenceClassification,
							        '''text-generation''': ElectraForCausalLM,
							        '''token-classification''': ElectraForTokenClassification,
							        '''zero-shot''': ElectraForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							_SCREAMING_SNAKE_CASE  :		List[str]			  =       True







							def        lowercase__							(    self   :	int   ,     __UpperCAmelCase   :	List[str]   ,     __UpperCAmelCase   :	Union[str, Any]   ,     __UpperCAmelCase   :	Union[str, Any]=False				)		->					Any:

										"""simple docstring"""



										UpperCamelCase_     				= super()._prepare_for_class(__UpperCAmelCase   ,     __UpperCAmelCase   ,     return_labels=__UpperCAmelCase				)

										if return_labels:
													if model_class in get_values(__UpperCAmelCase				):
																UpperCamelCase_     				= torch.zeros(
																    (self.model_tester.batch_size, self.model_tester.seq_length)   ,     dtype=torch.long   ,     device=__UpperCAmelCase				)
										return inputs_dict







							def        lowercase__							(    self   :	str				)		->					Tuple:

										"""simple docstring"""



										UpperCamelCase_     				= ElectraModelTester(self				)
										UpperCamelCase_     				= ConfigTester(self   ,     config_class=__UpperCAmelCase   ,     hidden_size=37				)







							def        lowercase__							(    self   :	List[str]				)		->					Union[str, Any]:

										"""simple docstring"""



										self.config_tester.run_common_tests()







							def        lowercase__							(    self   :	str				)		->					List[str]:

										"""simple docstring"""



										UpperCamelCase_     				= self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_model(*__UpperCAmelCase				)







							def        lowercase__							(    self   :	List[Any]				)		->					Dict:

										"""simple docstring"""



										UpperCamelCase_     				= self.model_tester.prepare_config_and_inputs_for_decoder()
										self.model_tester.create_and_check_electra_model_as_decoder(*__UpperCAmelCase				)







							def        lowercase__							(    self   :	int				)		->					Union[str, Any]:

										"""simple docstring"""



										UpperCamelCase_     				= self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
													UpperCamelCase_     				= type
													self.model_tester.create_and_check_electra_model(*__UpperCAmelCase				)







							def        lowercase__							(    self   :	List[Any]				)		->					Optional[int]:

										"""simple docstring"""



										UpperCamelCase_     				= self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_masked_lm(*__UpperCAmelCase				)







							def        lowercase__							(    self   :	Optional[Any]				)		->					Optional[Any]:

										"""simple docstring"""



										UpperCamelCase_     				= self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_token_classification(*__UpperCAmelCase				)







							def        lowercase__							(    self   :	Union[str, Any]				)		->					Union[str, Any]:

										"""simple docstring"""



										UpperCamelCase_     				= self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_pretraining(*__UpperCAmelCase				)







							def        lowercase__							(    self   :	int				)		->					Optional[int]:

										"""simple docstring"""



										UpperCamelCase_     				= self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_sequence_classification(*__UpperCAmelCase				)







							def        lowercase__							(    self   :	Optional[Any]				)		->					Optional[int]:

										"""simple docstring"""



										UpperCamelCase_     				= self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_question_answering(*__UpperCAmelCase				)







							def        lowercase__							(    self   :	Optional[int]				)		->					Tuple:

										"""simple docstring"""



										UpperCamelCase_     				= self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_electra_for_multiple_choice(*__UpperCAmelCase				)







							@slow
							def        lowercase__							(    self   :	int				)		->					Union[str, Any]:

										"""simple docstring"""



										for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													UpperCamelCase_     				= ElectraModel.from_pretrained(__UpperCAmelCase				)
													self.assertIsNotNone(__UpperCAmelCase				)







							def        lowercase__							(    self   :	str				)		->					Tuple:

										"""simple docstring"""



										UpperCamelCase_     				= self.model_tester.prepare_config_and_inputs_for_decoder()
										self.model_tester.create_and_check_electra_for_causal_lm(*__UpperCAmelCase				)







@require_torch
class     A					(				unittest.TestCase       ):







							@slow
							def        lowercase__							(    self   :	Union[str, Any]				)		->					int:

										"""simple docstring"""



										UpperCamelCase_     				= ElectraModel.from_pretrained('google/electra-small-discriminator'				)
										UpperCamelCase_     				= torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]				)
										UpperCamelCase_     				= torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]				)
										UpperCamelCase_     				= model(__UpperCAmelCase   ,     attention_mask=__UpperCAmelCase				)[0]
										UpperCamelCase_     				= torch.Size((1, 11, 256)				)
										self.assertEqual(output.shape   ,     __UpperCAmelCase				)
										UpperCamelCase_     				= torch.tensor(
										    [[[0.4_471, 0.6_821, -0.3_265], [0.4_627, 0.5_255, -0.3_668], [0.4_532, 0.3_313, -0.4_344]]]				)

										self.assertTrue(torch.allclose(output[:, 1:4, 1:4]   ,     __UpperCAmelCase   ,     atol=1E-4				)				)






from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			torch.Tensor:

			'''simple docstring'''


			if not inplace:
						UpperCamelCase_     				= ma + ma
			else:
						ma += ma

			return ma




def  a_		(							__snake_case     ,		__snake_case					)	->			torch.Tensor:

			'''simple docstring'''


			UpperCamelCase_     				= -1 * len(__snake_case					)
			UpperCamelCase_     				= list(range(len(tensor.shape[:zero_index]					)					)					)
			return tensor.permute(first_inds + [zero_index + i for i in inds]					)




def  a_		(							__snake_case     ,		__snake_case					)	->			torch.Tensor:

			'''simple docstring'''


			return t.reshape(t.shape[:-no_dims] + (-1,)					)




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case     ,		__snake_case = 1E-4					)	->			torch.Tensor:

			'''simple docstring'''


			UpperCamelCase_     				= mask.expand(*value.shape					)
			return torch.sum(mask * value     ,		dim=__snake_case					) / (eps + torch.sum(__snake_case     ,		dim=__snake_case					))




def  a_		(							__snake_case     ,		__snake_case = 2.31_25     ,		__snake_case = 21.68_75     ,		__snake_case = 6_4					)	->			torch.Tensor:

			'''simple docstring'''


			UpperCamelCase_     				= torch.linspace(__snake_case     ,		__snake_case     ,		no_bins - 1     ,		device=pts.device					)
			UpperCamelCase_     				= torch.sqrt(torch.sum((pts.unsqueeze(-2					) - pts.unsqueeze(-3					)) ** 2     ,		dim=-1					)					)
			return torch.bucketize(__snake_case     ,		__snake_case					)




def  a_		(							__snake_case     ,		__snake_case					)	->			dict:

			'''simple docstring'''


			UpperCamelCase_     				= dicts[0]
			UpperCamelCase_     				= {}
			for k, v in first.items():
						UpperCamelCase_     				= [d[k] for d in dicts]
						if isinstance(__snake_case     ,		__snake_case					):
									UpperCamelCase_     				= dict_multimap(__snake_case     ,		__snake_case					)
						else:
									UpperCamelCase_     				= fn(__snake_case					)

			return new_dict




def  a_		(							__snake_case     ,		__snake_case					)	->			torch.Tensor:

			'''simple docstring'''


			UpperCamelCase_     				= v_bins.view(((1,) * len(x.shape					)) + (len(__snake_case					),)					)
			UpperCamelCase_     				= x[..., None] - reshaped_bins
			UpperCamelCase_     				= torch.argmin(torch.abs(__snake_case					)     ,		dim=-1					)
			return nn.functional.one_hot(__snake_case     ,		num_classes=len(__snake_case					)					).float()




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case = 0     ,		__snake_case = 0					)	->			torch.Tensor:

			'''simple docstring'''


			UpperCamelCase_     				= []
			for i, s in enumerate(data.shape[:no_batch_dims]					):
						UpperCamelCase_     				= torch.arange(__snake_case					)
						UpperCamelCase_     				= r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape					) - i - 1)))					)
						ranges.append(__snake_case					)

			UpperCamelCase_     				= [slice(__snake_case					) for _ in range(len(data.shape					) - no_batch_dims					)]
			UpperCamelCase_     				= inds
			ranges.extend(__snake_case					)
			# Matt note: Editing this to get around the behaviour of using a list as an array index changing
			# in recent Numpy versions
			return data[tuple(__snake_case					)]


__a			:	List[Any]							     =		TypeVar("""T""")




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			Dict[Any, Union[dict, list, tuple, Any]]:

			'''simple docstring'''


			UpperCamelCase_     				= {}
			for k, v in dic.items():
						if isinstance(__snake_case     ,		__snake_case					):
									UpperCamelCase_     				= dict_map(__snake_case     ,		__snake_case     ,		__snake_case					)
						else:
									UpperCamelCase_     				= tree_map(__snake_case     ,		__snake_case     ,		__snake_case					)

			return new_dict




@overload
def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			Any:

			'''simple docstring'''


			...




@overload
def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			dict:

			'''simple docstring'''


			...




@overload
def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			list:

			'''simple docstring'''


			...




@overload
def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			tuple:

			'''simple docstring'''


			...




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case					)	->			Any:

			'''simple docstring'''


			if isinstance(__snake_case     ,		__snake_case					):
						return dict_map(__snake_case     ,		__snake_case     ,		__snake_case					)
			elif isinstance(__snake_case     ,		__snake_case					):
						return [tree_map(__snake_case     ,		__snake_case     ,		__snake_case					) for x in tree]
			elif isinstance(__snake_case     ,		__snake_case					):
						return tuple(tree_map(__snake_case     ,		__snake_case     ,		__snake_case					) for x in tree					)
			elif isinstance(__snake_case     ,		__snake_case					):
						return fn(__snake_case					)
			else:
						print(type(__snake_case					)					)
						raise ValueError('Not supported'					)





__a			:	Tuple							     =		partial(tree_map, leaf_type=torch.Tensor)
