



import os
import tempfile
import unittest

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

from ...generation.test_utils import GenerationTesterMixin
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,
		    NezhaForMaskedLM,
		    NezhaForMultipleChoice,
		    NezhaForNextSentencePrediction,
		    NezhaForPreTraining,
		    NezhaForQuestionAnswering,
		    NezhaForSequenceClassification,
		    NezhaForTokenClassification,
		    NezhaModel,
		)
		from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST





class 			__magic_name__ :






		def __init__(    self   :  List[Any]					,		UpperCamelCase__   :  Any					,		UpperCamelCase__   :  Union[str, Any]=13					,		UpperCamelCase__   :  Tuple=7					,		UpperCamelCase__   :  Tuple=True					,		UpperCamelCase__   :  Dict=True					,		UpperCamelCase__   :  Union[str, Any]=True					,		UpperCamelCase__   :  Union[str, Any]=True					,		UpperCamelCase__   :  str=99					,		UpperCamelCase__   :  List[Any]=32					,		UpperCamelCase__   :  Dict=5					,		UpperCamelCase__   :  List[str]=4					,		UpperCamelCase__   :  List[Any]=37					,		UpperCamelCase__   :  Any="gelu"					,		UpperCamelCase__   :  int=0.1					,		UpperCamelCase__   :  str=0.1					,		UpperCamelCase__   :  List[str]=1_28					,		UpperCamelCase__   :  Optional[Any]=32					,		UpperCamelCase__   :  int=16					,		UpperCamelCase__   :  int=2					,		UpperCamelCase__   :  Optional[int]=0.02					,		UpperCamelCase__   :  Optional[Any]=3					,		UpperCamelCase__   :  Tuple=4					,		UpperCamelCase__   :  List[str]=None					,		)					->       List[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 SCREAMING_SNAKE_CASE_      (    self   :  Optional[int]   )					->       Optional[Any]:
								'''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								=  self.get_config()

								return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels






		def SCREAMING_SNAKE_CASE_      (    self   :  List[Any]   )					->       Any:
								'''simple docstring'''







								return NezhaConfig(
								    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=__a					,		initializer_range=self.initializer_range					,		)






		def SCREAMING_SNAKE_CASE_      (    self   :  Union[str, Any]   )					->       Optional[int]:
								'''simple docstring'''







								(
								    (
								    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 SCREAMING_SNAKE_CASE_      (    self   :  Any					,		UpperCamelCase__   :  List[str]					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  Optional[Any]					,		UpperCamelCase__   :  Optional[Any]					,		UpperCamelCase__   :  Dict					,		UpperCamelCase__   :  Optional[Any]					,		UpperCamelCase__   :  str   )					->       Union[str, Any]:
								'''simple docstring'''







								UpperCAmelCase								=  NezhaModel(config=__a   )
								model.to(__a   )
								model.eval()
								UpperCAmelCase								=  model(__a					,		attention_mask=__a					,		token_type_ids=__a   )
								UpperCAmelCase								=  model(__a					,		token_type_ids=__a   )
								UpperCAmelCase								=  model(__a   )
								self.parent.assertEqual(result.last_hidden_state.shape					,		(self.batch_size, self.seq_length, self.hidden_size)   )
								self.parent.assertEqual(result.pooler_output.shape					,		(self.batch_size, self.hidden_size)   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Union[str, Any]					,		UpperCamelCase__   :  List[Any]					,		UpperCamelCase__   :  List[str]					,		UpperCamelCase__   :  Tuple					,		UpperCamelCase__   :  Optional[int]					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  List[Any]					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  Union[str, Any]					,		)					->       Dict:
								'''simple docstring'''







								UpperCAmelCase								=  True
								UpperCAmelCase								=  NezhaModel(__a   )
								model.to(__a   )
								model.eval()
								UpperCAmelCase								=  model(
								    __a					,		attention_mask=__a					,		token_type_ids=__a					,		encoder_hidden_states=__a					,		encoder_attention_mask=__a					,		)
								UpperCAmelCase								=  model(
								    __a					,		attention_mask=__a					,		token_type_ids=__a					,		encoder_hidden_states=__a					,		)
								UpperCAmelCase								=  model(__a					,		attention_mask=__a					,		token_type_ids=__a   )
								self.parent.assertEqual(result.last_hidden_state.shape					,		(self.batch_size, self.seq_length, self.hidden_size)   )
								self.parent.assertEqual(result.pooler_output.shape					,		(self.batch_size, self.hidden_size)   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Optional[Any]					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  List[str]					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  Any					,		UpperCamelCase__   :  str					,		UpperCamelCase__   :  str					,		UpperCamelCase__   :  Union[str, Any]   )					->       Dict:
								'''simple docstring'''







								UpperCAmelCase								=  NezhaForMaskedLM(config=__a   )
								model.to(__a   )
								model.eval()
								UpperCAmelCase								=  model(__a					,		attention_mask=__a					,		token_type_ids=__a					,		labels=__a   )
								self.parent.assertEqual(result.logits.shape					,		(self.batch_size, self.seq_length, self.vocab_size)   )






		def SCREAMING_SNAKE_CASE_      (    self   :  List[str]					,		UpperCamelCase__   :  Union[str, Any]					,		UpperCamelCase__   :  Union[str, Any]					,		UpperCamelCase__   :  List[Any]					,		UpperCamelCase__   :  List[Any]					,		UpperCamelCase__   :  Tuple					,		UpperCamelCase__   :  Union[str, Any]					,		UpperCamelCase__   :  Tuple   )					->       Any:
								'''simple docstring'''







								UpperCAmelCase								=  NezhaForNextSentencePrediction(config=__a   )
								model.to(__a   )
								model.eval()
								UpperCAmelCase								=  model(
								    __a					,		attention_mask=__a					,		token_type_ids=__a					,		labels=__a					,		)
								self.parent.assertEqual(result.logits.shape					,		(self.batch_size, 2)   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Any					,		UpperCamelCase__   :  List[str]					,		UpperCamelCase__   :  str					,		UpperCamelCase__   :  Any					,		UpperCamelCase__   :  Dict					,		UpperCamelCase__   :  List[Any]					,		UpperCamelCase__   :  Optional[Any]					,		UpperCamelCase__   :  str   )					->       Dict:
								'''simple docstring'''







								UpperCAmelCase								=  NezhaForPreTraining(config=__a   )
								model.to(__a   )
								model.eval()
								UpperCAmelCase								=  model(
								    __a					,		attention_mask=__a					,		token_type_ids=__a					,		labels=__a					,		next_sentence_label=__a					,		)
								self.parent.assertEqual(result.prediction_logits.shape					,		(self.batch_size, self.seq_length, self.vocab_size)   )
								self.parent.assertEqual(result.seq_relationship_logits.shape					,		(self.batch_size, 2)   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Union[str, Any]					,		UpperCamelCase__   :  Tuple					,		UpperCamelCase__   :  Tuple					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  Optional[int]					,		UpperCamelCase__   :  List[Any]					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  Tuple   )					->       List[Any]:
								'''simple docstring'''







								UpperCAmelCase								=  NezhaForQuestionAnswering(config=__a   )
								model.to(__a   )
								model.eval()
								UpperCAmelCase								=  model(
								    __a					,		attention_mask=__a					,		token_type_ids=__a					,		start_positions=__a					,		end_positions=__a					,		)
								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 SCREAMING_SNAKE_CASE_      (    self   :  List[Any]					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  Optional[Any]					,		UpperCamelCase__   :  Tuple					,		UpperCamelCase__   :  Optional[int]					,		UpperCamelCase__   :  Optional[int]					,		UpperCamelCase__   :  Optional[Any]					,		UpperCamelCase__   :  int   )					->       Dict:
								'''simple docstring'''







								UpperCAmelCase								=  self.num_labels
								UpperCAmelCase								=  NezhaForSequenceClassification(__a   )
								model.to(__a   )
								model.eval()
								UpperCAmelCase								=  model(__a					,		attention_mask=__a					,		token_type_ids=__a					,		labels=__a   )
								self.parent.assertEqual(result.logits.shape					,		(self.batch_size, self.num_labels)   )






		def SCREAMING_SNAKE_CASE_      (    self   :  List[str]					,		UpperCamelCase__   :  Union[str, Any]					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  List[str]					,		UpperCamelCase__   :  Optional[int]					,		UpperCamelCase__   :  Optional[Any]					,		UpperCamelCase__   :  Dict					,		UpperCamelCase__   :  Any   )					->       List[str]:
								'''simple docstring'''







								UpperCAmelCase								=  self.num_labels
								UpperCAmelCase								=  NezhaForTokenClassification(config=__a   )
								model.to(__a   )
								model.eval()
								UpperCAmelCase								=  model(__a					,		attention_mask=__a					,		token_type_ids=__a					,		labels=__a   )
								self.parent.assertEqual(result.logits.shape					,		(self.batch_size, self.seq_length, self.num_labels)   )






		def SCREAMING_SNAKE_CASE_      (    self   :  List[Any]					,		UpperCamelCase__   :  Optional[int]					,		UpperCamelCase__   :  Optional[Any]					,		UpperCamelCase__   :  Tuple					,		UpperCamelCase__   :  Dict					,		UpperCamelCase__   :  Dict					,		UpperCamelCase__   :  Any					,		UpperCamelCase__   :  Optional[Any]   )					->       Union[str, Any]:
								'''simple docstring'''







								UpperCAmelCase								=  self.num_choices
								UpperCAmelCase								=  NezhaForMultipleChoice(config=__a   )
								model.to(__a   )
								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(
								    __a					,		attention_mask=__a					,		token_type_ids=__a					,		labels=__a					,		)
								self.parent.assertEqual(result.logits.shape					,		(self.batch_size, self.num_choices)   )







		def SCREAMING_SNAKE_CASE_      (    self   :  int   )					->       Tuple:
								'''simple docstring'''







								UpperCAmelCase								=  self.prepare_config_and_inputs()
								(
								    (
								    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 			__magic_name__ (				A__,	A__,	A__,	unittest.TestCase    ):
		lowercase   :							Dict        =(
		    (
		        NezhaModel,
		        NezhaForMaskedLM,
		        NezhaForMultipleChoice,
		        NezhaForNextSentencePrediction,
		        NezhaForPreTraining,
		        NezhaForQuestionAnswering,
		        NezhaForSequenceClassification,
		        NezhaForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		lowercase   :							Dict        =(
		    {
		        '''feature-extraction''': NezhaModel,
		        '''fill-mask''': NezhaForMaskedLM,
		        '''question-answering''': NezhaForQuestionAnswering,
		        '''text-classification''': NezhaForSequenceClassification,
		        '''token-classification''': NezhaForTokenClassification,
		        '''zero-shot''': NezhaForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		lowercase   :							Any        =True






		def SCREAMING_SNAKE_CASE_      (    self   :  Dict					,		UpperCamelCase__   :  int					,		UpperCamelCase__   :  Optional[Any]					,		UpperCamelCase__   :  int=False   )					->       str:
								'''simple docstring'''







								UpperCAmelCase								=  super()._prepare_for_class(__a					,		__a					,		return_labels=__a   )

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






		def SCREAMING_SNAKE_CASE_      (    self   :  List[str]   )					->       int:
								'''simple docstring'''







								UpperCAmelCase								=  NezhaModelTester(self   )
								UpperCAmelCase								=  ConfigTester(self					,		config_class=__a					,		hidden_size=37   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Optional[Any]   )					->       Any:
								'''simple docstring'''







								self.config_tester.run_common_tests()






		def SCREAMING_SNAKE_CASE_      (    self   :  Dict   )					->       str:
								'''simple docstring'''







								UpperCAmelCase								=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*__a   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Tuple   )					->       Optional[int]:
								'''simple docstring'''







								UpperCAmelCase								=  self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(*__a   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Optional[int]   )					->       Any:
								'''simple docstring'''







								(
								    (
								    UpperCAmelCase    
								)							,		(
								    UpperCAmelCase    
								)							,		(
								    UpperCAmelCase    
								)							,		(
								    UpperCAmelCase    
								)							,		(
								    UpperCAmelCase    
								)							,		(
								    UpperCAmelCase    
								)							,		(
								    UpperCAmelCase    
								)							,		(
								    UpperCAmelCase    
								)							,		(
								    UpperCAmelCase    
								)							,		    
								)								=  self.model_tester.prepare_config_and_inputs_for_decoder()

								UpperCAmelCase								=  None

								self.model_tester.create_and_check_model_as_decoder(
								    __a					,		__a					,		__a					,		__a					,		__a					,		__a					,		__a					,		__a					,		__a					,		)






		def SCREAMING_SNAKE_CASE_      (    self   :  Optional[Any]   )					->       Any:
								'''simple docstring'''







								UpperCAmelCase								=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*__a   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Optional[int]   )					->       Optional[int]:
								'''simple docstring'''







								UpperCAmelCase								=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*__a   )






		def SCREAMING_SNAKE_CASE_      (    self   :  List[str]   )					->       Tuple:
								'''simple docstring'''







								UpperCAmelCase								=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_next_sequence_prediction(*__a   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Tuple   )					->       Tuple:
								'''simple docstring'''







								UpperCAmelCase								=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_pretraining(*__a   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Dict   )					->       str:
								'''simple docstring'''







								UpperCAmelCase								=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*__a   )






		def SCREAMING_SNAKE_CASE_      (    self   :  Optional[int]   )					->       Any:
								'''simple docstring'''







								UpperCAmelCase								=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*__a   )






		def SCREAMING_SNAKE_CASE_      (    self   :  int   )					->       List[str]:
								'''simple docstring'''







								UpperCAmelCase								=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*__a   )






		@slow
		def SCREAMING_SNAKE_CASE_      (    self   :  Union[str, Any]   )					->       Union[str, Any]:
								'''simple docstring'''







								for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														UpperCAmelCase								=  NezhaModel.from_pretrained(__a   )
														self.assertIsNotNone(__a   )







		@slow
		@require_torch_gpu
		def SCREAMING_SNAKE_CASE_      (    self   :  Tuple   )					->       Tuple:
								'''simple docstring'''







								UpperCAmelCase							,		UpperCAmelCase								=  self.model_tester.prepare_config_and_inputs_for_common()
								for model_class in self.all_model_classes:
														# NezhaForMultipleChoice behaves incorrectly in JIT environments.
														if model_class == NezhaForMultipleChoice:
																				return

														UpperCAmelCase								=  True
														UpperCAmelCase								=  model_class(config=__a   )

														UpperCAmelCase								=  self._prepare_for_class(__a					,		__a   )
														UpperCAmelCase								=  torch.jit.trace(
														    __a					,		(inputs_dict["input_ids"].to("cpu"   ), inputs_dict["attention_mask"].to("cpu"   ))   )

														with tempfile.TemporaryDirectory() as tmp:
																				torch.jit.save(__a					,		os.path.join(__a					,		"bert.pt"   )   )
																				UpperCAmelCase								=  torch.jit.load(os.path.join(__a					,		"bert.pt"   )					,		map_location=__a   )
																				loaded(inputs_dict["input_ids"].to(__a   )					,		inputs_dict["attention_mask"].to(__a   )   )





@require_torch
class 			__magic_name__ (				unittest.TestCase    ):






		@slow
		def SCREAMING_SNAKE_CASE_      (    self   :  Optional[Any]   )					->       Union[str, Any]:
								'''simple docstring'''







								UpperCAmelCase								=  NezhaModel.from_pretrained("sijunhe/nezha-cn-base"   )
								UpperCAmelCase								=  torch.tensor([[0, 1, 2, 3, 4, 5]]   )
								UpperCAmelCase								=  torch.tensor([[0, 1, 1, 1, 1, 1]]   )
								with torch.no_grad():
														UpperCAmelCase								=  model(__a					,		attention_mask=__a   )[0]
								UpperCAmelCase								=  torch.Size((1, 6, 7_68)   )
								self.assertEqual(output.shape					,		__a   )
								UpperCAmelCase								=  torch.tensor([[[0.06_85, 0.24_41, 0.11_02], [0.06_00, 0.19_06, 0.13_49], [0.02_21, 0.08_19, 0.05_86]]]   )

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







		@slow
		def SCREAMING_SNAKE_CASE_      (    self   :  Tuple   )					->       Tuple:
								'''simple docstring'''







								UpperCAmelCase								=  NezhaForMaskedLM.from_pretrained("sijunhe/nezha-cn-base"   )
								UpperCAmelCase								=  torch.tensor([[0, 1, 2, 3, 4, 5]]   )
								UpperCAmelCase								=  torch.tensor([[1, 1, 1, 1, 1, 1]]   )
								with torch.no_grad():
														UpperCAmelCase								=  model(__a					,		attention_mask=__a   )[0]
								UpperCAmelCase								=  torch.Size((1, 6, 2_11_28)   )
								self.assertEqual(output.shape					,		__a   )
								UpperCAmelCase								=  torch.tensor(
								    [[-2.79_39, -1.79_02, -2.21_89], [-2.85_85, -1.89_08, -2.37_23], [-2.64_99, -1.77_50, -2.25_58]]   )

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





"""simple docstring"""





from collections.abc import Callable

import numpy as np



def  lowerCamelCase__						(     __snake_case,	__snake_case,	__snake_case,	__snake_case,	__snake_case      )				->						np.array:


					"""simple docstring"""





					_UpperCamelCase									=    int(np.ceil((x_end - xa) / step_size      )      )
					_UpperCamelCase									=    np.zeros((n + 1,)      )
					_UpperCamelCase									=    ya
					_UpperCamelCase									=    xa

					for k in range(__snake_case      ):
										_UpperCamelCase									=    y[k] + step_size * ode_func(__snake_case,	y[k]      )
										_UpperCamelCase									=    y[k] + (
										    (step_size / 2) * (ode_func(__snake_case,	y[k]      ) + ode_func(x + step_size,	__snake_case      ))
										)
										x += step_size

					return y


if __name__ == "__main__":
				import doctest

				doctest.testmod()






