"""simple docstring"""






import unittest

from transformers import is_torch_available
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 (
          RemBertConfig,
          RemBertForCausalLM,
          RemBertForMaskedLM,
          RemBertForMultipleChoice,
          RemBertForQuestionAnswering,
          RemBertForSequenceClassification,
          RemBertForTokenClassification,
          RemBertModel,
      )
      from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST
class     snake_case	:




     """simple docstring"""

     def __init__(    self       :     List[str]   ,lowerCamelCase__       :     Dict   ,lowerCamelCase__       :     Optional[Any]=13   ,lowerCamelCase__       :     Any=7   ,lowerCamelCase__       :     Optional[Any]=True   ,lowerCamelCase__       :     Any=True   ,lowerCamelCase__       :     int=True   ,lowerCamelCase__       :     List[Any]=True   ,lowerCamelCase__       :     Optional[int]=99   ,lowerCamelCase__       :     Optional[int]=32   ,lowerCamelCase__       :     Optional[Any]=18   ,lowerCamelCase__       :     Optional[Any]=43   ,lowerCamelCase__       :     Union[str, Any]=5   ,lowerCamelCase__       :     Optional[Any]=4   ,lowerCamelCase__       :     List[str]=37   ,lowerCamelCase__       :     Optional[int]="gelu"   ,lowerCamelCase__       :     Union[str, Any]=0.1   ,lowerCamelCase__       :     Optional[Any]=0.1   ,lowerCamelCase__       :     Any=512   ,lowerCamelCase__       :     Optional[Any]=16   ,lowerCamelCase__       :     List[str]=2   ,lowerCamelCase__       :     List[Any]=0.0_2   ,lowerCamelCase__       :     Dict=3   ,lowerCamelCase__       :     List[str]=4   ,lowerCamelCase__       :     Optional[int]=None   ,):
         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__					=							input_embedding_size
         UpperCAmelCase__					=							output_embedding_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 		__lowerCAmelCase						(    self       :     int      ):
         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__					=							RemBertConfig(
             vocab_size=self.vocab_size   ,hidden_size=self.hidden_size   ,input_embedding_size=self.input_embedding_size   ,output_embedding_size=self.output_embedding_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=lowerCamelCase__   ,initializer_range=self.initializer_range   ,)

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







     def 		__lowerCAmelCase						(    self       :     Optional[Any]      ):
         (
             (
             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 		__lowerCAmelCase						(    self       :     Tuple   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     List[str]   ,lowerCamelCase__       :     List[str]   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     Dict   ,lowerCamelCase__       :     int      ):
         UpperCAmelCase__					=							RemBertModel(config=lowerCamelCase__      )
         model.to(lowerCamelCase__      )
         model.eval()
         UpperCAmelCase__					=							model(lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,token_type_ids=lowerCamelCase__      )
         UpperCAmelCase__					=							model(lowerCamelCase__   ,token_type_ids=lowerCamelCase__      )
         UpperCAmelCase__					=							model(lowerCamelCase__      )
         self.parent.assertEqual(result.last_hidden_state.shape   ,(self.batch_size, self.seq_length, self.hidden_size)      )







     def 		__lowerCAmelCase						(    self       :     List[str]   ,lowerCamelCase__       :     List[Any]   ,lowerCamelCase__       :     Any   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     Dict   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     Dict   ,lowerCamelCase__       :     List[Any]   ,):
         UpperCAmelCase__					=							True
         UpperCAmelCase__					=							RemBertModel(lowerCamelCase__      )
         model.to(lowerCamelCase__      )
         model.eval()
         UpperCAmelCase__					=							model(
             lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,token_type_ids=lowerCamelCase__   ,encoder_hidden_states=lowerCamelCase__   ,encoder_attention_mask=lowerCamelCase__   ,)
         UpperCAmelCase__					=							model(
             lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,token_type_ids=lowerCamelCase__   ,encoder_hidden_states=lowerCamelCase__   ,)
         UpperCAmelCase__					=							model(lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,token_type_ids=lowerCamelCase__      )
         self.parent.assertEqual(result.last_hidden_state.shape   ,(self.batch_size, self.seq_length, self.hidden_size)      )







     def 		__lowerCAmelCase						(    self       :     List[str]   ,lowerCamelCase__       :     Dict   ,lowerCamelCase__       :     Any   ,lowerCamelCase__       :     Dict   ,lowerCamelCase__       :     int   ,lowerCamelCase__       :     Dict   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     Dict   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     Union[str, Any]   ,):
         UpperCAmelCase__					=							RemBertForCausalLM(config=lowerCamelCase__      )
         model.to(lowerCamelCase__      )
         model.eval()
         UpperCAmelCase__					=							model(lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,token_type_ids=lowerCamelCase__   ,labels=lowerCamelCase__      )
         self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.vocab_size)      )







     def 		__lowerCAmelCase						(    self       :     Optional[int]   ,lowerCamelCase__       :     Any   ,lowerCamelCase__       :     List[str]   ,lowerCamelCase__       :     str   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     int   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     Optional[int]      ):
         UpperCAmelCase__					=							RemBertForMaskedLM(config=lowerCamelCase__      )
         model.to(lowerCamelCase__      )
         model.eval()
         UpperCAmelCase__					=							model(lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,token_type_ids=lowerCamelCase__   ,labels=lowerCamelCase__      )
         self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.vocab_size)      )







     def 		__lowerCAmelCase						(    self       :     Optional[int]   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     List[str]   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     Any   ,lowerCamelCase__       :     Dict   ,lowerCamelCase__       :     List[Any]   ,lowerCamelCase__       :     List[str]   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     Union[str, Any]   ,):
         UpperCAmelCase__					=							True
         UpperCAmelCase__					=							True
         UpperCAmelCase__					=							RemBertForCausalLM(config=lowerCamelCase__      )
         model.to(lowerCamelCase__      )
         model.eval()

         # first forward pass
         UpperCAmelCase__					=							model(
             lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,encoder_hidden_states=lowerCamelCase__   ,encoder_attention_mask=lowerCamelCase__   ,use_cache=lowerCamelCase__   ,)
         UpperCAmelCase__					=							outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         UpperCAmelCase__					=							ids_tensor((self.batch_size, 3)   ,config.vocab_size      )
         UpperCAmelCase__					=							ids_tensor((self.batch_size, 3)   ,vocab_size=2      )

         # append to next input_ids and
         UpperCAmelCase__					=							torch.cat([input_ids, next_tokens]   ,dim=-1      )
         UpperCAmelCase__					=							torch.cat([input_mask, next_mask]   ,dim=-1      )

         UpperCAmelCase__					=							model(
             lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,encoder_hidden_states=lowerCamelCase__   ,encoder_attention_mask=lowerCamelCase__   ,output_hidden_states=lowerCamelCase__   ,)['hidden_states'][0]
         UpperCAmelCase__					=							model(
             lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,encoder_hidden_states=lowerCamelCase__   ,encoder_attention_mask=lowerCamelCase__   ,past_key_values=lowerCamelCase__   ,output_hidden_states=lowerCamelCase__   ,)['hidden_states'][0]

         # select random slice
         UpperCAmelCase__					=							ids_tensor((1,)   ,output_from_past.shape[-1]      ).item()
         UpperCAmelCase__					=							output_from_no_past[:, -3:, random_slice_idx].detach()
         UpperCAmelCase__					=							output_from_past[:, :, random_slice_idx].detach()

         self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]      )

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(lowerCamelCase__   ,lowerCamelCase__   ,atol=1e-3      )      )







     def 		__lowerCAmelCase						(    self       :     Dict   ,lowerCamelCase__       :     List[Any]   ,lowerCamelCase__       :     Optional[int]   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     List[Any]   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     Dict      ):
         UpperCAmelCase__					=							RemBertForQuestionAnswering(config=lowerCamelCase__      )
         model.to(lowerCamelCase__      )
         model.eval()
         UpperCAmelCase__					=							model(
             lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,token_type_ids=lowerCamelCase__   ,start_positions=lowerCamelCase__   ,end_positions=lowerCamelCase__   ,)
         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       :     Optional[int]   ,lowerCamelCase__       :     Optional[int]   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     Any   ,lowerCamelCase__       :     Union[str, Any]      ):
         UpperCAmelCase__					=							self.num_labels
         UpperCAmelCase__					=							RemBertForSequenceClassification(lowerCamelCase__      )
         model.to(lowerCamelCase__      )
         model.eval()
         UpperCAmelCase__					=							model(lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,token_type_ids=lowerCamelCase__   ,labels=lowerCamelCase__      )
         self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.num_labels)      )







     def 		__lowerCAmelCase						(    self       :     Tuple   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     str   ,lowerCamelCase__       :     Dict   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     Tuple      ):
         UpperCAmelCase__					=							self.num_labels
         UpperCAmelCase__					=							RemBertForTokenClassification(config=lowerCamelCase__      )
         model.to(lowerCamelCase__      )
         model.eval()
         UpperCAmelCase__					=							model(lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,token_type_ids=lowerCamelCase__   ,labels=lowerCamelCase__      )
         self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.num_labels)      )







     def 		__lowerCAmelCase						(    self       :     Any   ,lowerCamelCase__       :     List[str]   ,lowerCamelCase__       :     Optional[int]   ,lowerCamelCase__       :     Any   ,lowerCamelCase__       :     Union[str, Any]   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     Optional[int]   ,lowerCamelCase__       :     List[str]      ):
         UpperCAmelCase__					=							self.num_choices
         UpperCAmelCase__					=							RemBertForMultipleChoice(config=lowerCamelCase__      )
         model.to(lowerCamelCase__      )
         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(
             lowerCamelCase__   ,attention_mask=lowerCamelCase__   ,token_type_ids=lowerCamelCase__   ,labels=lowerCamelCase__   ,)
         self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.num_choices)      )







     def 		__lowerCAmelCase						(    self       :     Any      ):
         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     snake_case	(			__UpperCAmelCase				,   __UpperCAmelCase				,   unittest.TestCase    ):




     """simple docstring"""

     snake_case__       =			(
         (
             RemBertModel,
             RemBertForMaskedLM,
             RemBertForCausalLM,
             RemBertForMultipleChoice,
             RemBertForQuestionAnswering,
             RemBertForSequenceClassification,
             RemBertForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     snake_case__       =			(RemBertForCausalLM,) if is_torch_available() else ()
     snake_case__       =			(
         {
             "feature-extraction": RemBertModel,
             "fill-mask": RemBertForMaskedLM,
             "question-answering": RemBertForQuestionAnswering,
             "text-classification": RemBertForSequenceClassification,
             "text-generation": RemBertForCausalLM,
             "token-classification": RemBertForTokenClassification,
             "zero-shot": RemBertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )







     def 		__lowerCAmelCase						(    self       :     Dict      ):
         UpperCAmelCase__					=							RemBertModelTester(self      )
         UpperCAmelCase__					=							ConfigTester(self   ,config_class=lowerCamelCase__   ,hidden_size=37      )







     def 		__lowerCAmelCase						(    self       :     List[str]      ):
         self.config_tester.run_common_tests()







     def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
         UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     Union[str, Any]      ):
         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_model(*lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     Dict      ):
         UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     int      ):
         UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     Optional[Any]      ):
         UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(*lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     Any      ):
         UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     int      ):
         UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     Optional[Any]      ):
         UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
         UpperCAmelCase__					=							self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(*lowerCamelCase__      )







     def 		__lowerCAmelCase						(    self       :     Tuple      ):
         # This regression test was failing with PyTorch < 1.3
         (
             (
             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(
             lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__   ,)







     @slow
     def 		__lowerCAmelCase						(    self       :     Any      ):
         for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             UpperCAmelCase__					=							RemBertModel.from_pretrained(lowerCamelCase__      )
             self.assertIsNotNone(lowerCamelCase__      )






@require_torch
class     snake_case	(			unittest.TestCase    ):




     """simple docstring"""

     @slow
     def 		__lowerCAmelCase						(    self       :     Dict      ):
         # Test exact values at the last hidden layer
         UpperCAmelCase__					=							RemBertModel.from_pretrained('google/rembert'      )
         UpperCAmelCase__					=							torch.tensor([[312, 56_498, 313, 2_125, 313]]      )
         UpperCAmelCase__					=							torch.tensor([[0, 0, 0, 1, 1]]      )
         with torch.no_grad():
             UpperCAmelCase__					=							model(lowerCamelCase__   ,token_type_ids=lowerCamelCase__   ,output_hidden_states=lowerCamelCase__      )

         UpperCAmelCase__					=							1_152

         UpperCAmelCase__					=							torch.Size((1, 5, hidden_size)      )
         self.assertEqual(output['last_hidden_state'].shape   ,lowerCamelCase__      )

         UpperCAmelCase__					=							torch.tensor(
             [
                 [
                     [0.0_7_5_4, -0.2_0_2_2, 0.1_9_0_4],
                     [-0.3_3_5_4, -0.3_6_9_2, -0.4_7_9_1],
                     [-0.2_3_1_4, -0.6_7_2_9, -0.0_7_4_9],
                     [-0.0_3_9_6, -0.3_1_0_5, -0.4_2_3_4],
                     [-0.1_5_7_1, -0.0_5_2_5, 0.5_3_5_3],
                 ]
             ]      )

         # Running on the original tf implementation gives slightly different results here.
         # Not clear why this variations is present
         # TODO: Find reason for discrepancy
         # expected_original_implementation = [[
         #     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
         #     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
         #     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
         #     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
         #     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
         # ]]

         self.assertTrue(torch.allclose(output['last_hidden_state'][:, :, :3]   ,lowerCamelCase__   ,atol=1e-4      )      )




"""simple docstring"""






from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve



def   a_						( lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase	):
				UpperCAmelCase__					=							np.stack((horizontal_flow, vertical_flow) ,      2	)

				# Create a grid of all pixel coordinates and subtract the flow to get the
				# target pixels coordinates
				UpperCAmelCase__					=							np.stack(
				    np.meshgrid(np.arange(0 ,      image.shape[1]	) ,      np.arange(0 ,      image.shape[0]	)	) ,      2	)
				UpperCAmelCase__					=							np.round(grid - flow	).astype(np.intaa	)

				# Find the locations outside of the original image
				UpperCAmelCase__					=							(grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]	))
				UpperCAmelCase__					=							0

				UpperCAmelCase__					=							image[grid[:, :, 1], grid[:, :, 0]]

				# Set pixels at invalid locations to 0
				UpperCAmelCase__					=							0

				return warped



def   a_						( lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase = None ,      ):
				if alpha is None:
								UpperCAmelCase__					=							0.1

				# Initialize flow
				UpperCAmelCase__					=							np.zeros_like(lowerCamelCase	)
				UpperCAmelCase__					=							np.zeros_like(lowerCamelCase	)

				# Prepare kernels for the calculation of the derivatives and the average velocity
				UpperCAmelCase__					=							np.array([[-1, 1], [-1, 1]]	) * 0.25
				UpperCAmelCase__					=							np.array([[-1, -1], [1, 1]]	) * 0.25
				UpperCAmelCase__					=							np.array([[1, 1], [1, 1]]	) * 0.25
				UpperCAmelCase__					=							np.array(
				    [[1 / 1_2, 1 / 6, 1 / 1_2], [1 / 6, 0, 1 / 6], [1 / 1_2, 1 / 6, 1 / 1_2]]	)

				# Iteratively refine the flow
				for _ in range(lowerCamelCase	):
								UpperCAmelCase__					=							warp(lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase	)
								UpperCAmelCase__					=							convolve(lowerCamelCase ,      lowerCamelCase	) + convolve(lowerCamelCase ,      lowerCamelCase	)
								UpperCAmelCase__					=							convolve(lowerCamelCase ,      lowerCamelCase	) + convolve(lowerCamelCase ,      lowerCamelCase	)
								UpperCAmelCase__					=							convolve(lowerCamelCase ,      lowerCamelCase	) + convolve(lowerCamelCase ,      -kernel_t	)

								UpperCAmelCase__					=							convolve(lowerCamelCase ,      lowerCamelCase	)
								UpperCAmelCase__					=							convolve(lowerCamelCase ,      lowerCamelCase	)

								# This updates the flow as proposed in the paper (Step 12)
								UpperCAmelCase__					=							(
								    derivative_x * avg_horizontal_velocity
								    + derivative_y * avg_vertical_velocity
								    + derivative_t
								)
								UpperCAmelCase__					=							update / (alpha**2 + derivative_x**2 + derivative_y**2)

								UpperCAmelCase__					=							avg_horizontal_velocity - derivative_x * update
								UpperCAmelCase__					=							avg_vertical_velocity - derivative_y * update

				return horizontal_flow, vertical_flow


if __name__ == "__main__":
						import doctest

						doctest.testmod()



