



"""simple docstring"""







from __future__ import annotations

import copy
import inspect
import json
import math
import os
import tempfile
import unittest
from importlib import import_module

import numpy as np

from transformers import ViTMAEConfig
from transformers.file_utils import cached_property, is_tf_available, is_vision_available
from transformers.testing_utils import require_tf, require_vision, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers import TFViTMAEForPreTraining, TFViTMAEModel


if is_vision_available():
     from PIL import Image

     from transformers import ViTImageProcessor







class 				UpperCamelCase__       :




      """simple docstring"""




      def __init__(				self    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Any=1_3   ,      SCREAMING_SNAKE_CASE_    :				str=3_0   ,      SCREAMING_SNAKE_CASE_    :				str=2   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=3   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=True   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=True   ,      SCREAMING_SNAKE_CASE_    :				Any=3_2   ,      SCREAMING_SNAKE_CASE_    :				Tuple=2   ,      SCREAMING_SNAKE_CASE_    :				Dict=4   ,      SCREAMING_SNAKE_CASE_    :				Any=3_7   ,      SCREAMING_SNAKE_CASE_    :				str="gelu"   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=0.1   ,      SCREAMING_SNAKE_CASE_    :				int=0.1   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=1_0   ,      SCREAMING_SNAKE_CASE_    :				str=0.02   ,      SCREAMING_SNAKE_CASE_    :				Any=3   ,      SCREAMING_SNAKE_CASE_    :				int=0.6   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=None   ,      ):
        lowerCAmelCase_       :    Optional[Any]								=			parent
        lowerCAmelCase_       :    Union[str, Any]								=			batch_size
        lowerCAmelCase_       :    Union[str, Any]								=			image_size
        lowerCAmelCase_       :    int								=			patch_size
        lowerCAmelCase_       :    Any								=			num_channels
        lowerCAmelCase_       :    Union[str, Any]								=			is_training
        lowerCAmelCase_       :    Union[str, Any]								=			use_labels
        lowerCAmelCase_       :    Optional[int]								=			hidden_size
        lowerCAmelCase_       :    int								=			num_hidden_layers
        lowerCAmelCase_       :    List[Any]								=			num_attention_heads
        lowerCAmelCase_       :    Optional[int]								=			intermediate_size
        lowerCAmelCase_       :    Dict								=			hidden_act
        lowerCAmelCase_       :    List[Any]								=			hidden_dropout_prob
        lowerCAmelCase_       :    Tuple								=			attention_probs_dropout_prob
        lowerCAmelCase_       :    Optional[int]								=			type_sequence_label_size
        lowerCAmelCase_       :    Dict								=			initializer_range
        lowerCAmelCase_       :    List[Any]								=			mask_ratio
        lowerCAmelCase_       :    List[Any]								=			scope

        # in ViTMAE, the expected sequence length = (num_patches + 1) * (1 - config.mask_ratio), rounded above
        # (we add 1 for the [CLS] token)
        lowerCAmelCase_       :    int								=			(image_size // patch_size) ** 2
        lowerCAmelCase_       :    int								=			int(math.ceil((1 - mask_ratio) * (num_patches + 1)	)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        lowerCAmelCase_       :    List[str]								=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

        lowerCAmelCase_       :    Union[str, Any]								=			None
        if self.use_labels:
          lowerCAmelCase_       :    Dict								=			ids_tensor([self.batch_size]   ,      self.type_sequence_label_size	)

        lowerCAmelCase_       :    List[Any]								=			self.get_config()

        return config, pixel_values, labels



      def 					SCREAMING_SNAKE_CASE__			(				self    :				int	):
        return ViTMAEConfig(
            image_size=self.image_size   ,      patch_size=self.patch_size   ,      num_channels=self.num_channels   ,      hidden_size=self.hidden_size   ,      num_hidden_layers=self.num_hidden_layers   ,      num_attention_heads=self.num_attention_heads   ,      intermediate_size=self.intermediate_size   ,      decoder_hidden_size=self.hidden_size   ,      decoder_num_hidden_layers=self.num_hidden_layers   ,      decoder_num_attention_heads=self.num_attention_heads   ,      decoder_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   ,      is_decoder=SCREAMING_SNAKE_CASE_   ,      initializer_range=self.initializer_range   ,      mask_ratio=self.mask_ratio   ,      )



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]	):
        lowerCAmelCase_       :    Any								=			TFViTMAEModel(config=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    str								=			model(SCREAMING_SNAKE_CASE_   ,      training=SCREAMING_SNAKE_CASE_	)
        self.parent.assertEqual(result.last_hidden_state.shape   ,      (self.batch_size, self.seq_length, self.hidden_size)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]	):
        lowerCAmelCase_       :    Optional[int]								=			TFViTMAEForPreTraining(SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    str								=			model(SCREAMING_SNAKE_CASE_   ,      training=SCREAMING_SNAKE_CASE_	)
        # expected sequence length = num_patches
        lowerCAmelCase_       :    List[Any]								=			(self.image_size // self.patch_size) ** 2
        lowerCAmelCase_       :    str								=			self.patch_size**2 * self.num_channels
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, num_patches, expected_num_channels)	)

        # test greyscale images
        lowerCAmelCase_       :    List[str]								=			1
        lowerCAmelCase_       :    str								=			TFViTMAEForPreTraining(SCREAMING_SNAKE_CASE_	)

        lowerCAmelCase_       :    Optional[Any]								=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
        lowerCAmelCase_       :    str								=			model(SCREAMING_SNAKE_CASE_   ,      training=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Any								=			self.patch_size**2
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, num_patches, expected_num_channels)	)





      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple	):
        lowerCAmelCase_       :    Dict								=			self.prepare_config_and_inputs()
        ((lowerCAmelCase_)						,(lowerCAmelCase_)						,(lowerCAmelCase_))       :    List[Any]								=			config_and_inputs
        lowerCAmelCase_       :    Any								=			{'pixel_values': pixel_values}
        return config, inputs_dict







@require_tf
class 				UpperCamelCase__       (     lowercase_, lowercase_, unittest.TestCase  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE      =					(TFViTMAEModel, TFViTMAEForPreTraining) if is_tf_available() else ()
      _SCREAMING_SNAKE_CASE      =					{"""feature-extraction""": TFViTMAEModel} if is_tf_available() else {}

      _SCREAMING_SNAKE_CASE      =					False
      _SCREAMING_SNAKE_CASE      =					False
      _SCREAMING_SNAKE_CASE      =					False
      _SCREAMING_SNAKE_CASE      =					False



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Dict	):
        lowerCAmelCase_       :    Union[str, Any]								=			TFViTMAEModelTester(self	)
        lowerCAmelCase_       :    Optional[Any]								=			ConfigTester(self   ,      config_class=SCREAMING_SNAKE_CASE_   ,      has_text_modality=SCREAMING_SNAKE_CASE_   ,      hidden_size=3_7	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Dict	):
        self.config_tester.run_common_tests()



      @unittest.skip(reason='ViTMAE does not use inputs_embeds'	)
      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        pass



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        lowerCAmelCase_						,lowerCAmelCase_       :    List[Any]								=			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          lowerCAmelCase_       :    Tuple								=			model_class(SCREAMING_SNAKE_CASE_	)
          self.assertIsInstance(model.get_input_embeddings()   ,      (tf.keras.layers.Layer)	)
          lowerCAmelCase_       :    Tuple								=			model.get_output_embeddings()
          self.assertTrue(x is None or isinstance(SCREAMING_SNAKE_CASE_   ,      tf.keras.layers.Layer	)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple	):
        lowerCAmelCase_						,lowerCAmelCase_       :    Union[str, Any]								=			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          lowerCAmelCase_       :    List[str]								=			model_class(SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    int								=			inspect.signature(model.call	)
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          lowerCAmelCase_       :    Dict								=			[*signature.parameters.keys()]

          lowerCAmelCase_       :    Optional[int]								=			['pixel_values']
          self.assertListEqual(arg_names[:1]   ,      SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):
        lowerCAmelCase_       :    List[Any]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        lowerCAmelCase_       :    int								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				str	):
        # make the mask reproducible
        np.random.seed(2	)

        lowerCAmelCase_						,lowerCAmelCase_       :    int								=			self.model_tester.prepare_config_and_inputs_for_common()

        lowerCAmelCase_       :    List[str]								=			int((config.image_size // config.patch_size) ** 2	)
        lowerCAmelCase_       :    List[Any]								=			np.random.uniform(size=(self.model_tester.batch_size, num_patches)	)

        for model_class in self.all_model_classes:
          lowerCAmelCase_       :    List[Any]								=			model_class(SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    Union[str, Any]								=			self._prepare_for_class(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)

          lowerCAmelCase_       :    Optional[int]								=			model(SCREAMING_SNAKE_CASE_   ,      noise=SCREAMING_SNAKE_CASE_	)

          lowerCAmelCase_       :    Tuple								=			copy.deepcopy(self._prepare_for_class(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)	)
          lowerCAmelCase_       :    Optional[int]								=			model(**SCREAMING_SNAKE_CASE_   ,      noise=SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    Any								=			outputs_dict[0].numpy()
          lowerCAmelCase_       :    Any								=			outputs_keywords[0].numpy()

          self.assertLess(np.sum(np.abs(output_dict - output_keywords	)	)   ,      1E-6	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        # make the mask reproducible
        np.random.seed(2	)

        lowerCAmelCase_						,lowerCAmelCase_       :    Union[str, Any]								=			self.model_tester.prepare_config_and_inputs_for_common()

        lowerCAmelCase_       :    Tuple								=			int((config.image_size // config.patch_size) ** 2	)
        lowerCAmelCase_       :    List[Any]								=			np.random.uniform(size=(self.model_tester.batch_size, num_patches)	)

        def prepare_numpy_arrays(SCREAMING_SNAKE_CASE_    :				List[str]	):
          lowerCAmelCase_       :    Union[str, Any]								=			{}
          for k, v in inputs_dict.items():
            if tf.is_tensor(SCREAMING_SNAKE_CASE_	):
              lowerCAmelCase_       :    str								=			v.numpy()
            else:
              lowerCAmelCase_       :    List[Any]								=			np.array(SCREAMING_SNAKE_CASE_	)

          return inputs_np_dict

        for model_class in self.all_model_classes:
          lowerCAmelCase_       :    Optional[int]								=			model_class(SCREAMING_SNAKE_CASE_	)

          lowerCAmelCase_       :    Any								=			self._prepare_for_class(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    Any								=			prepare_numpy_arrays(SCREAMING_SNAKE_CASE_	)

          lowerCAmelCase_       :    List[str]								=			model(SCREAMING_SNAKE_CASE_   ,      noise=SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    List[str]								=			model(**SCREAMING_SNAKE_CASE_   ,      noise=SCREAMING_SNAKE_CASE_	)
          self.assert_outputs_same(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Any	):
        # make masks reproducible
        np.random.seed(2	)

        lowerCAmelCase_       :    Tuple								=			int((tf_model.config.image_size // tf_model.config.patch_size) ** 2	)
        lowerCAmelCase_       :    Dict								=			np.random.uniform(size=(self.model_tester.batch_size, num_patches)	)
        lowerCAmelCase_       :    Optional[Any]								=			tf.constant(SCREAMING_SNAKE_CASE_	)

        # Add `noise` argument.
        # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
        lowerCAmelCase_       :    int								=			tf_noise

        super().check_pt_tf_models(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]	):
        # make mask reproducible
        np.random.seed(2	)

        lowerCAmelCase_						,lowerCAmelCase_       :    int								=			self.model_tester.prepare_config_and_inputs_for_common()

        lowerCAmelCase_       :    Tuple								=			{
            module_member
            for model_class in self.all_model_classes
            for module in (import_module(model_class.__module__	),)
            for module_member_name in dir(SCREAMING_SNAKE_CASE_	)
            if module_member_name.endswith('MainLayer'	)
            # This condition is required, since `modeling_tf_clip.py` has 3 classes whose names end with `MainLayer`.
            and module_member_name[: -len('MainLayer'	)] == model_class.__name__[: -len('Model'	)]
            for module_member in (getattr(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	),)
            if isinstance(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
            and tf.keras.layers.Layer in module_member.__bases__
            and getattr(SCREAMING_SNAKE_CASE_   ,      '_keras_serializable'   ,      SCREAMING_SNAKE_CASE_	)
        }

        lowerCAmelCase_       :    int								=			int((config.image_size // config.patch_size) ** 2	)
        lowerCAmelCase_       :    Optional[int]								=			np.random.uniform(size=(self.model_tester.batch_size, num_patches)	)
        lowerCAmelCase_       :    List[Any]								=			tf.convert_to_tensor(SCREAMING_SNAKE_CASE_	)
        inputs_dict.update({'noise': noise}	)

        for main_layer_class in tf_main_layer_classes:
          lowerCAmelCase_       :    int								=			main_layer_class(SCREAMING_SNAKE_CASE_	)

          lowerCAmelCase_       :    Union[str, Any]								=			{
              name: tf.keras.Input(tensor.shape[1:]   ,      dtype=tensor.dtype	) for name, tensor in inputs_dict.items()
          }

          lowerCAmelCase_       :    int								=			tf.keras.Model(SCREAMING_SNAKE_CASE_   ,      outputs=main_layer(SCREAMING_SNAKE_CASE_	)	)
          lowerCAmelCase_       :    Any								=			model(SCREAMING_SNAKE_CASE_	)

          with tempfile.TemporaryDirectory() as tmpdirname:
            lowerCAmelCase_       :    int								=			os.path.join(SCREAMING_SNAKE_CASE_   ,      'keras_model.h5'	)
            model.save(SCREAMING_SNAKE_CASE_	)
            lowerCAmelCase_       :    Tuple								=			tf.keras.models.load_model(
                SCREAMING_SNAKE_CASE_   ,      custom_objects={main_layer_class.__name__: main_layer_class}	)
            assert isinstance(SCREAMING_SNAKE_CASE_   ,      tf.keras.Model	)
            lowerCAmelCase_       :    int								=			model(SCREAMING_SNAKE_CASE_	)
            self.assert_outputs_same(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)



      @slow
      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        # make mask reproducible
        np.random.seed(2	)

        lowerCAmelCase_						,lowerCAmelCase_       :    Any								=			self.model_tester.prepare_config_and_inputs_for_common()

        lowerCAmelCase_       :    List[str]								=			int((config.image_size // config.patch_size) ** 2	)
        lowerCAmelCase_       :    List[str]								=			np.random.uniform(size=(self.model_tester.batch_size, num_patches)	)

        for model_class in self.all_model_classes:
          lowerCAmelCase_       :    int								=			model_class(SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    int								=			self._prepare_for_class(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    List[str]								=			model(SCREAMING_SNAKE_CASE_   ,      noise=SCREAMING_SNAKE_CASE_	)

          if model_class.__name__ == "TFViTMAEModel":
            lowerCAmelCase_       :    Optional[Any]								=			outputs.last_hidden_state.numpy()
            lowerCAmelCase_       :    Optional[Any]								=			0
          else:
            lowerCAmelCase_       :    List[Any]								=			outputs.logits.numpy()
            lowerCAmelCase_       :    List[str]								=			0

          with tempfile.TemporaryDirectory() as tmpdirname:
            model.save_pretrained(SCREAMING_SNAKE_CASE_   ,      saved_model=SCREAMING_SNAKE_CASE_	)
            lowerCAmelCase_       :    Optional[int]								=			model_class.from_pretrained(SCREAMING_SNAKE_CASE_	)
            lowerCAmelCase_       :    Union[str, Any]								=			model(SCREAMING_SNAKE_CASE_   ,      noise=SCREAMING_SNAKE_CASE_	)

            if model_class.__name__ == "TFViTMAEModel":
              lowerCAmelCase_       :    List[str]								=			after_outputs['last_hidden_state'].numpy()
              lowerCAmelCase_       :    int								=			0
            else:
              lowerCAmelCase_       :    List[str]								=			after_outputs['logits'].numpy()
              lowerCAmelCase_       :    List[str]								=			0

            lowerCAmelCase_       :    List[str]								=			np.amax(np.abs(out_a - out_a	)	)
            self.assertLessEqual(SCREAMING_SNAKE_CASE_   ,      1E-5	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]	):
        # make mask reproducible
        np.random.seed(2	)

        lowerCAmelCase_						,lowerCAmelCase_       :    Dict								=			self.model_tester.prepare_config_and_inputs_for_common()

        lowerCAmelCase_       :    List[str]								=			int((config.image_size // config.patch_size) ** 2	)
        lowerCAmelCase_       :    Optional[int]								=			np.random.uniform(size=(self.model_tester.batch_size, num_patches)	)

        for model_class in self.all_model_classes:
          lowerCAmelCase_       :    Any								=			model_class(SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    List[Any]								=			self._prepare_for_class(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)

          lowerCAmelCase_       :    str								=			model(SCREAMING_SNAKE_CASE_   ,      noise=SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    Tuple								=			model.get_config()
          # make sure that returned config is jsonifiable, which is required by keras
          json.dumps(SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    List[str]								=			model_class.from_config(model.get_config()	)
          # make sure it also accepts a normal config
          lowerCAmelCase_       :    Union[str, Any]								=			model_class.from_config(model.config	)
          lowerCAmelCase_       :    Optional[Any]								=			new_model(SCREAMING_SNAKE_CASE_	)  # Build model
          new_model.set_weights(model.get_weights()	)
          lowerCAmelCase_       :    str								=			new_model(SCREAMING_SNAKE_CASE_   ,      noise=SCREAMING_SNAKE_CASE_	)

          self.assert_outputs_same(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)



      @unittest.skip(
          reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results.'	)
      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]	):
        pass



      @unittest.skip(reason='ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load'	)
      def 					SCREAMING_SNAKE_CASE__			(				self    :				str	):
        pass





      @slow
      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        lowerCAmelCase_       :    Any								=			TFViTMAEModel.from_pretrained('google/vit-base-patch16-224'	)
        self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)




def     UpperCamelCase_   (   )						->       Optional[Any]:

  """simple docstring"""




  lowerCAmelCase_       :    Union[str, Any]								=			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'						)
  return image







@require_tf
@require_vision
class 				UpperCamelCase__       (     unittest.TestCase  ):




      """simple docstring"""




      @cached_property
      def 					SCREAMING_SNAKE_CASE__			(				self    :				int	):
        return ViTImageProcessor.from_pretrained('facebook/vit-mae-base'	) if is_vision_available() else None





      @slow
      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):
        # make random mask reproducible across the PT and TF model
        np.random.seed(2	)

        lowerCAmelCase_       :    str								=			TFViTMAEForPreTraining.from_pretrained('facebook/vit-mae-base'	)

        lowerCAmelCase_       :    Optional[int]								=			self.default_image_processor
        lowerCAmelCase_       :    Optional[Any]								=			prepare_img()
        lowerCAmelCase_       :    Union[str, Any]								=			image_processor(images=SCREAMING_SNAKE_CASE_   ,      return_tensors='tf'	)

        # prepare a noise vector that will be also used for testing the TF model
        # (this way we can ensure that the PT and TF models operate on the same inputs)
        lowerCAmelCase_       :    Dict								=			ViTMAEConfig()
        lowerCAmelCase_       :    List[Any]								=			int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2	)
        lowerCAmelCase_       :    Tuple								=			np.random.uniform(size=(1, num_patches)	)

        # forward pass
        lowerCAmelCase_       :    int								=			model(**SCREAMING_SNAKE_CASE_   ,      noise=SCREAMING_SNAKE_CASE_	)

        # verify the logits
        lowerCAmelCase_       :    Any								=			tf.convert_to_tensor([1, 1_9_6, 7_6_8]	)
        self.assertEqual(outputs.logits.shape   ,      SCREAMING_SNAKE_CASE_	)

        lowerCAmelCase_       :    List[Any]								=			tf.convert_to_tensor(
            [[-0.05_48, -1.70_23, -0.93_25], [0.37_21, -0.56_70, -0.22_33], [0.82_35, -1.38_78, -0.35_24]]	)

        tf.debugging.assert_near(outputs.logits[0, :3, :3]   ,      SCREAMING_SNAKE_CASE_   ,      atol=1E-4	)










"""simple docstring"""







import unittest

from transformers import AlbertConfig, 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, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         MODEL_FOR_PRETRAINING_MAPPING,
         AlbertForMaskedLM,
         AlbertForMultipleChoice,
         AlbertForPreTraining,
         AlbertForQuestionAnswering,
         AlbertForSequenceClassification,
         AlbertForTokenClassification,
         AlbertModel,
     )
     from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST







class 				UpperCamelCase__       :




      """simple docstring"""




      def __init__(				self    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict=1_3   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=7   ,      SCREAMING_SNAKE_CASE_    :				Dict=True   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=True   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=True   ,      SCREAMING_SNAKE_CASE_    :				str=True   ,      SCREAMING_SNAKE_CASE_    :				List[str]=9_9   ,      SCREAMING_SNAKE_CASE_    :				int=1_6   ,      SCREAMING_SNAKE_CASE_    :				List[str]=3_6   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=6   ,      SCREAMING_SNAKE_CASE_    :				Tuple=6   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=6   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=3_7   ,      SCREAMING_SNAKE_CASE_    :				Tuple="gelu"   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=0.1   ,      SCREAMING_SNAKE_CASE_    :				int=0.1   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=5_1_2   ,      SCREAMING_SNAKE_CASE_    :				List[str]=1_6   ,      SCREAMING_SNAKE_CASE_    :				List[str]=2   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=0.02   ,      SCREAMING_SNAKE_CASE_    :				Dict=3   ,      SCREAMING_SNAKE_CASE_    :				int=4   ,      SCREAMING_SNAKE_CASE_    :				Tuple=None   ,      ):
        lowerCAmelCase_       :    Any								=			parent
        lowerCAmelCase_       :    Optional[int]								=			batch_size
        lowerCAmelCase_       :    Dict								=			seq_length
        lowerCAmelCase_       :    Tuple								=			is_training
        lowerCAmelCase_       :    str								=			use_input_mask
        lowerCAmelCase_       :    Union[str, Any]								=			use_token_type_ids
        lowerCAmelCase_       :    Tuple								=			use_labels
        lowerCAmelCase_       :    Optional[int]								=			vocab_size
        lowerCAmelCase_       :    Any								=			embedding_size
        lowerCAmelCase_       :    Optional[Any]								=			hidden_size
        lowerCAmelCase_       :    str								=			num_hidden_layers
        lowerCAmelCase_       :    Optional[Any]								=			num_hidden_groups
        lowerCAmelCase_       :    Dict								=			num_attention_heads
        lowerCAmelCase_       :    Optional[Any]								=			intermediate_size
        lowerCAmelCase_       :    Any								=			hidden_act
        lowerCAmelCase_       :    Union[str, Any]								=			hidden_dropout_prob
        lowerCAmelCase_       :    int								=			attention_probs_dropout_prob
        lowerCAmelCase_       :    int								=			max_position_embeddings
        lowerCAmelCase_       :    List[Any]								=			type_vocab_size
        lowerCAmelCase_       :    Any								=			type_sequence_label_size
        lowerCAmelCase_       :    Optional[int]								=			initializer_range
        lowerCAmelCase_       :    Tuple								=			num_labels
        lowerCAmelCase_       :    Dict								=			num_choices
        lowerCAmelCase_       :    Tuple								=			scope



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]	):
        lowerCAmelCase_       :    List[str]								=			ids_tensor([self.batch_size, self.seq_length]   ,      self.vocab_size	)

        lowerCAmelCase_       :    str								=			None
        if self.use_input_mask:
          lowerCAmelCase_       :    int								=			random_attention_mask([self.batch_size, self.seq_length]	)

        lowerCAmelCase_       :    List[Any]								=			None
        if self.use_token_type_ids:
          lowerCAmelCase_       :    str								=			ids_tensor([self.batch_size, self.seq_length]   ,      self.type_vocab_size	)

        lowerCAmelCase_       :    List[str]								=			None
        lowerCAmelCase_       :    Union[str, Any]								=			None
        lowerCAmelCase_       :    str								=			None
        if self.use_labels:
          lowerCAmelCase_       :    Any								=			ids_tensor([self.batch_size]   ,      self.type_sequence_label_size	)
          lowerCAmelCase_       :    List[str]								=			ids_tensor([self.batch_size, self.seq_length]   ,      self.num_labels	)
          lowerCAmelCase_       :    Dict								=			ids_tensor([self.batch_size]   ,      self.num_choices	)

        lowerCAmelCase_       :    List[Any]								=			self.get_config()

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



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Dict	):
        return AlbertConfig(
            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   ,      initializer_range=self.initializer_range   ,      num_hidden_groups=self.num_hidden_groups   ,      )



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				str	):
        lowerCAmelCase_       :    Union[str, Any]								=			AlbertModel(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    List[str]								=			model(SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    List[Any]								=			model(SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    int								=			model(SCREAMING_SNAKE_CASE_	)
        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    :				int   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Tuple	):
        lowerCAmelCase_       :    Optional[Any]								=			AlbertForPreTraining(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    Optional[Any]								=			model(
            SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_   ,      labels=SCREAMING_SNAKE_CASE_   ,      sentence_order_label=SCREAMING_SNAKE_CASE_   ,      )
        self.parent.assertEqual(result.prediction_logits.shape   ,      (self.batch_size, self.seq_length, self.vocab_size)	)
        self.parent.assertEqual(result.sop_logits.shape   ,      (self.batch_size, config.num_labels)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				Dict	):
        lowerCAmelCase_       :    str								=			AlbertForMaskedLM(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    str								=			model(SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_   ,      labels=SCREAMING_SNAKE_CASE_	)
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.vocab_size)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				str   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]	):
        lowerCAmelCase_       :    List[str]								=			AlbertForQuestionAnswering(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    Any								=			model(
            SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_   ,      start_positions=SCREAMING_SNAKE_CASE_   ,      end_positions=SCREAMING_SNAKE_CASE_   ,      )
        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    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict	):
        lowerCAmelCase_       :    Union[str, Any]								=			self.num_labels
        lowerCAmelCase_       :    Union[str, Any]								=			AlbertForSequenceClassification(SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    Union[str, Any]								=			model(SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_   ,      labels=SCREAMING_SNAKE_CASE_	)
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_labels)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				List[str]	):
        lowerCAmelCase_       :    List[str]								=			self.num_labels
        lowerCAmelCase_       :    List[Any]								=			AlbertForTokenClassification(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    str								=			model(SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_   ,      labels=SCREAMING_SNAKE_CASE_	)
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.num_labels)	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				str   ,      SCREAMING_SNAKE_CASE_    :				List[str]	):
        lowerCAmelCase_       :    Optional[Any]								=			self.num_choices
        lowerCAmelCase_       :    int								=			AlbertForMultipleChoice(config=SCREAMING_SNAKE_CASE_	)
        model.to(SCREAMING_SNAKE_CASE_	)
        model.eval()
        lowerCAmelCase_       :    Optional[int]								=			input_ids.unsqueeze(1	).expand(-1   ,      self.num_choices   ,      -1	).contiguous()
        lowerCAmelCase_       :    List[Any]								=			token_type_ids.unsqueeze(1	).expand(-1   ,      self.num_choices   ,      -1	).contiguous()
        lowerCAmelCase_       :    List[Any]								=			input_mask.unsqueeze(1	).expand(-1   ,      self.num_choices   ,      -1	).contiguous()
        lowerCAmelCase_       :    List[Any]								=			model(
            SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_   ,      token_type_ids=SCREAMING_SNAKE_CASE_   ,      labels=SCREAMING_SNAKE_CASE_   ,      )
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_choices)	)





      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]	):
        lowerCAmelCase_       :    Optional[int]								=			self.prepare_config_and_inputs()
        (
            (
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,(
            lowerCAmelCase_
        )						,
        )       :    Optional[int]								=			config_and_inputs
        lowerCAmelCase_       :    Dict								=			{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
        return config, inputs_dict







@require_torch
class 				UpperCamelCase__       (     lowercase_, lowercase_, unittest.TestCase  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE      =					(
          (
              AlbertModel,
              AlbertForPreTraining,
              AlbertForMaskedLM,
              AlbertForMultipleChoice,
              AlbertForSequenceClassification,
              AlbertForTokenClassification,
              AlbertForQuestionAnswering,
          )
          if is_torch_available()
          else ()
      )
      _SCREAMING_SNAKE_CASE      =					(
          {
              """feature-extraction""": AlbertModel,
              """fill-mask""": AlbertForMaskedLM,
              """question-answering""": AlbertForQuestionAnswering,
              """text-classification""": AlbertForSequenceClassification,
              """token-classification""": AlbertForTokenClassification,
              """zero-shot""": AlbertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _SCREAMING_SNAKE_CASE      =					True



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Dict   ,      SCREAMING_SNAKE_CASE_    :				List[str]   ,      SCREAMING_SNAKE_CASE_    :				str=False	):
        lowerCAmelCase_       :    List[str]								=			super()._prepare_for_class(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      return_labels=SCREAMING_SNAKE_CASE_	)

        if return_labels:
          if model_class in get_values(SCREAMING_SNAKE_CASE_	):
            lowerCAmelCase_       :    List[str]								=			torch.zeros(
                (self.model_tester.batch_size, self.model_tester.seq_length)   ,      dtype=torch.long   ,      device=SCREAMING_SNAKE_CASE_	)
            lowerCAmelCase_       :    Any								=			torch.zeros(
                self.model_tester.batch_size   ,      dtype=torch.long   ,      device=SCREAMING_SNAKE_CASE_	)
        return inputs_dict



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):
        lowerCAmelCase_       :    str								=			AlbertModelTester(self	)
        lowerCAmelCase_       :    Optional[int]								=			ConfigTester(self   ,      config_class=SCREAMING_SNAKE_CASE_   ,      hidden_size=3_7	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        self.config_tester.run_common_tests()



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple	):
        lowerCAmelCase_       :    List[Any]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        lowerCAmelCase_       :    Any								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple	):
        lowerCAmelCase_       :    Tuple								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):
        lowerCAmelCase_       :    List[Any]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):
        lowerCAmelCase_       :    str								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				str	):
        lowerCAmelCase_       :    Optional[int]								=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        lowerCAmelCase_       :    Tuple								=			self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
          lowerCAmelCase_       :    int								=			type
          self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)





      @slow
      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowerCAmelCase_       :    Optional[Any]								=			AlbertModel.from_pretrained(SCREAMING_SNAKE_CASE_	)
          self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)







@require_torch
class 				UpperCamelCase__       (     unittest.TestCase  ):




      """simple docstring"""




      @slow
      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]	):
        lowerCAmelCase_       :    Any								=			AlbertModel.from_pretrained('albert-base-v2'	)
        lowerCAmelCase_       :    Tuple								=			torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]	)
        lowerCAmelCase_       :    Tuple								=			torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]	)
        with torch.no_grad():
          lowerCAmelCase_       :    List[Any]								=			model(SCREAMING_SNAKE_CASE_   ,      attention_mask=SCREAMING_SNAKE_CASE_	)[0]
        lowerCAmelCase_       :    str								=			torch.Size((1, 1_1, 7_6_8)	)
        self.assertEqual(output.shape   ,      SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    int								=			torch.tensor(
            [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]]	)

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





