




"""simple docstring"""


import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
					from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
					from tensorflow.keras.optimizers import Adam






class       lowercase(  tf.keras.optimizers.schedules.LearningRateSchedule				):







	'''simple docstring'''





	def __init__(       self:	int,   a_:	Optional[int],   a_:	Optional[Any],   a_:	Union[str, Any],   a_:	List[Any] = 1.0,   a_:	Union[str, Any] = None,   ):




			'''simple docstring'''



			super().__init__()
			_snake_case							:	List[str]			=						initial_learning_rate
			_snake_case							:	Any			=						warmup_steps
			_snake_case							:	List[Any]			=						power
			_snake_case							:	List[str]			=						decay_schedule_fn
			_snake_case							:	Optional[int]			=						name


	def __call__(       self:	Any,   a_:	Union[str, Any]	):




			'''simple docstring'''



			with tf.name_scope(self.name or """WarmUp"""	) as name:
					# Implements polynomial warmup. i.e., if global_step < warmup_steps, the
					# learning rate will be `global_step/num_warmup_steps * init_lr`.
					_snake_case							:	List[Any]			=						tf.cast(lowerCAmelCase_,   tf.floataa	)
					_snake_case							:	Optional[int]			=						tf.cast(self.warmup_steps,   tf.floataa	)
					_snake_case							:	str			=						global_step_float / warmup_steps_float
					_snake_case							:	Tuple			=						self.initial_learning_rate * tf.math.pow(lowerCAmelCase_,   self.power	)
					return tf.cond(
					    global_step_float < warmup_steps_float,   lambda: warmup_learning_rate,   lambda: self.decay_schedule_fn(step - self.warmup_steps	),   name=lowerCAmelCase_,   )


	def 				UpperCamelCase_     (       self:	str	):




			'''simple docstring'''



			return {
			    "initial_learning_rate": self.initial_learning_rate,
			    "decay_schedule_fn": self.decay_schedule_fn,
			    "warmup_steps": self.warmup_steps,
			    "power": self.power,
			    "name": self.name,
			}







def       UpperCAmelCase__ (snake_case__							:      float	,     snake_case__							:      int	,     snake_case__							:      int	,     snake_case__							:      float = 0.0	,     snake_case__							:      float = 0.9	,     snake_case__							:      float = 0.9_99	,     snake_case__							:      float = 1e-8	,     snake_case__							:      Optional[float] = None	,     snake_case__							:      Optional[float] = None	,     snake_case__							:      float = 0.0	,     snake_case__							:      float = 1.0	,     snake_case__							:      Optional[List[str]] = None	,     ):
		"""simple docstring"""

		_snake_case							:	str			=						tf.keras.optimizers.schedules.PolynomialDecay(
		    initial_learning_rate=snake_case__	,     decay_steps=num_train_steps - num_warmup_steps	,     end_learning_rate=init_lr * min_lr_ratio	,     power=snake_case__	,     )
		if num_warmup_steps:
				_snake_case							:	Tuple			=						WarmUp(
				    initial_learning_rate=snake_case__	,     decay_schedule_fn=snake_case__	,     warmup_steps=snake_case__	,     )
		if weight_decay_rate > 0.0:
				_snake_case							:	Tuple			=						AdamWeightDecay(
				    learning_rate=snake_case__	,     weight_decay_rate=snake_case__	,     beta_a=snake_case__	,     beta_a=snake_case__	,     epsilon=snake_case__	,     clipnorm=snake_case__	,     global_clipnorm=snake_case__	,     exclude_from_weight_decay=["""LayerNorm""", """layer_norm""", """bias"""]	,     include_in_weight_decay=snake_case__	,     )
		else:
				_snake_case							:	List[str]			=						tf.keras.optimizers.Adam(
				    learning_rate=snake_case__	,     beta_a=snake_case__	,     beta_a=snake_case__	,     epsilon=snake_case__	,     clipnorm=snake_case__	,     global_clipnorm=snake_case__	,     )
		# We return the optimizer and the LR scheduler in order to better track the
		# evolution of the LR independently of the optimizer.
		return optimizer, lr_schedule






class       lowercase(  __lowerCAmelCase				):







	'''simple docstring'''





	def __init__(       self:	Optional[int],   a_:	Optional[Any] = 0.001,   a_:	Optional[int] = 0.9,   a_:	Dict = 0.999,   a_:	Union[str, Any] = 1E-7,   a_:	Tuple = False,   a_:	Dict = 0.0,   a_:	str = None,   a_:	Dict = None,   a_:	Optional[int] = "AdamWeightDecay",   **a_:	int,   ):




			'''simple docstring'''



			super().__init__(lowerCAmelCase_,   lowerCAmelCase_,   lowerCAmelCase_,   lowerCAmelCase_,   lowerCAmelCase_,   lowerCAmelCase_,   **lowerCAmelCase_	)
			_snake_case							:	List[str]			=						weight_decay_rate
			_snake_case							:	int			=						include_in_weight_decay
			_snake_case							:	Optional[int]			=						exclude_from_weight_decay


	@classmethod
	def 				UpperCamelCase_     (       cls:	int,   a_:	List[str]	):




			'''simple docstring'''



			_snake_case							:	Union[str, Any]			=						{"""WarmUp""": WarmUp}
			return super(lowerCAmelCase_,   cls	).from_config(lowerCAmelCase_,   custom_objects=lowerCAmelCase_	)


	def 				UpperCamelCase_     (       self:	Optional[Any],   a_:	List[Any],   a_:	int,   a_:	str	):




			'''simple docstring'''



			super(lowerCAmelCase_,   self	)._prepare_local(lowerCAmelCase_,   lowerCAmelCase_,   lowerCAmelCase_	)
			_snake_case							:	Optional[int]			=						tf.constant(
			    self.weight_decay_rate,   name="""adam_weight_decay_rate"""	)


	def 				UpperCamelCase_     (       self:	str,   a_:	Optional[Any],   a_:	int,   a_:	str	):




			'''simple docstring'''



			_snake_case							:	str			=						self._do_use_weight_decay(var.name	)
			if do_decay:
					return var.assign_sub(
					    learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]["""weight_decay_rate"""],   use_locking=self._use_locking,   )
			return tf.no_op()


	def 				UpperCamelCase_     (       self:	int,   a_:	Optional[int],   a_:	Union[str, Any]=None,   **a_:	int	):




			'''simple docstring'''



			_snake_case	,      _snake_case							:	str			=						list(zip(*lowerCAmelCase_	)	)
			return super(lowerCAmelCase_,   self	).apply_gradients(zip(lowerCAmelCase_,   lowerCAmelCase_	),   name=lowerCAmelCase_,   **lowerCAmelCase_	)


	def 				UpperCamelCase_     (       self:	Dict,   a_:	int,   a_:	List[Any],   a_:	List[str]	):




			'''simple docstring'''



			if apply_state is None:
					return self._decayed_lr_t[var_dtype], {}

			_snake_case							:	Optional[Any]			=						apply_state or {}
			_snake_case							:	str			=						apply_state.get((var_device, var_dtype)	)
			if coefficients is None:
					_snake_case							:	List[str]			=						self._fallback_apply_state(lowerCAmelCase_,   lowerCAmelCase_	)
					_snake_case							:	Tuple			=						coefficients

			return coefficients["lr_t"], {"apply_state": apply_state}


	def 				UpperCamelCase_     (       self:	Dict,   a_:	Tuple,   a_:	Union[str, Any],   a_:	Optional[Any]=None	):




			'''simple docstring'''



			_snake_case	,      _snake_case							:	Optional[Any]			=						self._get_lr(var.device,   var.dtype.base_dtype,   lowerCAmelCase_	)
			_snake_case							:	Dict			=						self._decay_weights_op(lowerCAmelCase_,   lowerCAmelCase_,   lowerCAmelCase_	)
			with tf.control_dependencies([decay]	):
					return super(lowerCAmelCase_,   self	)._resource_apply_dense(lowerCAmelCase_,   lowerCAmelCase_,   **lowerCAmelCase_	)


	def 				UpperCamelCase_     (       self:	List[str],   a_:	Optional[int],   a_:	Optional[Any],   a_:	Any,   a_:	int=None	):




			'''simple docstring'''



			_snake_case	,      _snake_case							:	Tuple			=						self._get_lr(var.device,   var.dtype.base_dtype,   lowerCAmelCase_	)
			_snake_case							:	int			=						self._decay_weights_op(lowerCAmelCase_,   lowerCAmelCase_,   lowerCAmelCase_	)
			with tf.control_dependencies([decay]	):
					return super(lowerCAmelCase_,   self	)._resource_apply_sparse(lowerCAmelCase_,   lowerCAmelCase_,   lowerCAmelCase_,   **lowerCAmelCase_	)


	def 				UpperCamelCase_     (       self:	int	):




			'''simple docstring'''



			_snake_case							:	List[str]			=						super().get_config()
			config.update({"""weight_decay_rate""": self.weight_decay_rate}	)
			return config


	def 				UpperCamelCase_     (       self:	Union[str, Any],   a_:	List[str]	):




			'''simple docstring'''



			if self.weight_decay_rate == 0:
					return False

			if self._include_in_weight_decay:
					for r in self._include_in_weight_decay:
							if re.search(lowerCAmelCase_,   lowerCAmelCase_	) is not None:
									return True

			if self._exclude_from_weight_decay:
					for r in self._exclude_from_weight_decay:
							if re.search(lowerCAmelCase_,   lowerCAmelCase_	) is not None:
									return False
			return True







class       lowercase(  __lowerCAmelCase				):







	'''simple docstring'''





	def __init__(       self:	Any	):




			'''simple docstring'''



			_snake_case							:	List[str]			=						[]
			_snake_case							:	Tuple			=						None


	@property
	def 				UpperCamelCase_     (       self:	int	):




			'''simple docstring'''



			if self._accum_steps is None:
					_snake_case							:	Any			=						tf.Variable(
					    tf.constant(0,   dtype=tf.intaa	),   trainable=lowerCAmelCase_,   synchronization=tf.VariableSynchronization.ON_READ,   aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,   )

			return self._accum_steps.value()


	@property
	def 				UpperCamelCase_     (       self:	List[str]	):




			'''simple docstring'''



			if not self._gradients:
					raise ValueError("""The accumulator should be called first to initialize the gradients"""	)
			return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]


	def __call__(       self:	Any,   a_:	Optional[Any]	):




			'''simple docstring'''



			if not self._gradients:
					_snake_case							:	Any			=						self.step  # Create the step variable.
					self._gradients.extend(
					    [
					        tf.Variable(
					            tf.zeros_like(lowerCAmelCase_	),   trainable=lowerCAmelCase_,   synchronization=tf.VariableSynchronization.ON_READ,   aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,   )
					        if gradient is not None
					        else gradient
					        for gradient in gradients
					    ]	)
			if len(lowerCAmelCase_	) != len(self._gradients	):
					raise ValueError(f"Expected {len(self._gradients	)} gradients, but got {len(lowerCAmelCase_	)}"	)

			for accum_gradient, gradient in zip(self._gradients,   lowerCAmelCase_	):
					if accum_gradient is not None and gradient is not None:
							accum_gradient.assign_add(lowerCAmelCase_	)

			self._accum_steps.assign_add(1	)


	def 				UpperCamelCase_     (       self:	Optional[Any]	):




			'''simple docstring'''



			if not self._gradients:
					return
			self._accum_steps.assign(0	)
			for gradient in self._gradients:
					if gradient is not None:
							gradient.assign(tf.zeros_like(lowerCAmelCase_	)	)








"""simple docstring"""


import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

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


if is_torch_available():
     import torch

     from transformers import UperNetForSemanticSegmentation
     from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor






class       lowercase:







 '''simple docstring'''





 def __init__(       self:	Dict,   a_:	Union[str, Any],   a_:	Tuple=13,   a_:	Dict=32,   a_:	Optional[Any]=3,   a_:	Optional[Any]=4,   a_:	Optional[int]=[10, 20, 30, 40],   a_:	Any=[2, 2, 3, 2],   a_:	Dict=True,   a_:	Dict=True,   a_:	List[str]=37,   a_:	Dict="gelu",   a_:	List[str]=10,   a_:	Union[str, Any]=0.02,   a_:	Any=["stage2", "stage3", "stage4"],   a_:	Optional[int]=3,   a_:	Tuple=None,   ):




   '''simple docstring'''



   _snake_case							:	Dict			=						parent
   _snake_case							:	Dict			=						batch_size
   _snake_case							:	Optional[Any]			=						image_size
   _snake_case							:	int			=						num_channels
   _snake_case							:	Tuple			=						num_stages
   _snake_case							:	int			=						hidden_sizes
   _snake_case							:	List[str]			=						depths
   _snake_case							:	str			=						is_training
   _snake_case							:	Dict			=						use_labels
   _snake_case							:	List[str]			=						intermediate_size
   _snake_case							:	Optional[int]			=						hidden_act
   _snake_case							:	Any			=						type_sequence_label_size
   _snake_case							:	List[str]			=						initializer_range
   _snake_case							:	Union[str, Any]			=						out_features
   _snake_case							:	Dict			=						num_labels
   _snake_case							:	int			=						scope
   _snake_case							:	Dict			=						num_stages


 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   _snake_case							:	Optional[Any]			=						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

   _snake_case							:	Optional[int]			=						None
   if self.use_labels:
     _snake_case							:	Optional[Any]			=						ids_tensor([self.batch_size],   self.type_sequence_label_size	)

   _snake_case							:	Tuple			=						self.get_config()

   return config, pixel_values, labels


 def 				UpperCamelCase_     (       self:	Tuple	):




   '''simple docstring'''



   return ConvNextConfig(
       num_channels=self.num_channels,   num_stages=self.num_stages,   hidden_sizes=self.hidden_sizes,   depths=self.depths,   is_training=self.is_training,   intermediate_size=self.intermediate_size,   hidden_act=self.hidden_act,   out_features=self.out_features,   )


 def 				UpperCamelCase_     (       self:	List[Any]	):




   '''simple docstring'''



   return UperNetConfig(
       backbone_config=self.get_backbone_config(),   hidden_size=512,   pool_scales=[1, 2, 3, 6],   use_auxiliary_head=a_,   auxiliary_loss_weight=0.4,   auxiliary_in_channels=40,   auxiliary_channels=256,   auxiliary_num_convs=1,   auxiliary_concat_input=a_,   loss_ignore_index=255,   num_labels=self.num_labels,   )


 def 				UpperCamelCase_     (       self:	Tuple,   a_:	List[Any],   a_:	Dict,   a_:	Tuple	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						UperNetForSemanticSegmentation(config=a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	Tuple			=						model(a_	)
   self.parent.assertEqual(
       result.logits.shape,   (self.batch_size, self.num_labels, self.image_size, self.image_size)	)


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	Any			=						self.prepare_config_and_inputs()
   (
       (
       _snake_case
   )	,      (
       _snake_case
   )	,      (
       _snake_case
   )	,      
   )							:	List[Any]			=						config_and_inputs
   _snake_case							:	Any			=						{"""pixel_values""": pixel_values}
   return config, inputs_dict







@require_torch
class       lowercase(  __a ,    __a ,    unittest.TestCase				):







 '''simple docstring'''





 lowercase__				      =      (UperNetForSemanticSegmentation,) if is_torch_available() else ()
 lowercase__				      =      {"image-segmentation": UperNetForSemanticSegmentation} if is_torch_available() else {}
 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False
 lowercase__				      =      False


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	List[str]			=						UperNetModelTester(self	)
   _snake_case							:	Dict			=						ConfigTester(self,   config_class=a_,   has_text_modality=a_,   hidden_size=37	)


 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   self.create_and_test_config_common_properties()
   self.config_tester.create_and_test_config_to_json_string()
   self.config_tester.create_and_test_config_to_json_file()
   self.config_tester.create_and_test_config_from_and_save_pretrained()
   self.config_tester.create_and_test_config_with_num_labels()
   self.config_tester.check_config_can_be_init_without_params()
   self.config_tester.check_config_arguments_init()


 def 				UpperCamelCase_     (       self:	List[Any]	):




   '''simple docstring'''



   return


 def 				UpperCamelCase_     (       self:	Union[str, Any]	):




   '''simple docstring'''



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

   for model_class in self.all_model_classes:
     _snake_case							:	Dict			=						model_class(a_	)
     _snake_case							:	Optional[int]			=						inspect.signature(model.forward	)
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     _snake_case							:	Tuple			=						[*signature.parameters.keys()]

     _snake_case							:	Any			=						["""pixel_values"""]
     self.assertListEqual(arg_names[:1],   a_	)


 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   _snake_case							:	Union[str, Any]			=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_semantic_segmentation(*a_	)


 @unittest.skip(reason="""UperNet does not use inputs_embeds"""	)
 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   pass


 @unittest.skip(reason="""UperNet does not support input and output embeddings"""	)
 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   pass


 @unittest.skip(reason="""UperNet does not have a base model"""	)
 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   pass


 @unittest.skip(reason="""UperNet does not have a base model"""	)
 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   pass


 @require_torch_multi_gpu
 @unittest.skip(reason="""UperNet has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""	)
 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   pass


 @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""	)
 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   pass


 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



   def check_hidden_states_output(a_:	Dict,   a_:	List[str],   a_:	Optional[int]	):
     _snake_case							:	Optional[Any]			=						model_class(a_	)
     model.to(a_	)
     model.eval()

     with torch.no_grad():
       _snake_case							:	Any			=						model(**self._prepare_for_class(a_,   a_	)	)

     _snake_case							:	Any			=						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

     _snake_case							:	List[str]			=						self.model_tester.num_stages
     self.assertEqual(len(a_	),   expected_num_stages + 1	)

     # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
     self.assertListEqual(
         list(hidden_states[0].shape[-2:]	),   [self.model_tester.image_size // 4, self.model_tester.image_size // 4],   )

   _snake_case	,      _snake_case							:	str			=						self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     _snake_case							:	int			=						True
     check_hidden_states_output(a_,   a_,   a_	)

     # check that output_hidden_states also work using config
     del inputs_dict["output_hidden_states"]
     _snake_case							:	Optional[int]			=						True

     check_hidden_states_output(a_,   a_,   a_	)


 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



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

   _snake_case							:	Tuple			=						_config_zero_init(a_	)
   _snake_case							:	Dict			=						_config_zero_init(configs_no_init.backbone_config	)
   for model_class in self.all_model_classes:
     _snake_case							:	Optional[int]			=						model_class(config=a_	)
     for name, param in model.named_parameters():
       if param.requires_grad:
         self.assertIn(
             ((param.data.mean() * 1E9).round() / 1E9).item(),   [0.0, 1.0],   msg=f"Parameter {name} of model {model_class} seems not properly initialized",   )


 @unittest.skip(reason="""UperNet does not have tied weights"""	)
 def 				UpperCamelCase_     (       self:	Union[str, Any]	):




   '''simple docstring'''



   pass


 @slow
 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case							:	int			=						UperNetForSemanticSegmentation.from_pretrained(a_	)
     self.assertIsNotNone(a_	)







def       UpperCAmelCase__ ():
  """simple docstring"""

  _snake_case							:	Union[str, Any]			=						hf_hub_download(
      repo_id="""hf-internal-testing/fixtures_ade20k"""	,     repo_type="""dataset"""	,     filename="""ADE_val_00000001.jpg"""					)
  _snake_case							:	List[Any]			=						Image.open(snake_case__					).convert("""RGB"""					)
  return image






@require_torch
@require_vision
@slow
class       lowercase(  unittest.TestCase				):







 '''simple docstring'''





 def 				UpperCamelCase_     (       self:	Tuple	):




   '''simple docstring'''



   _snake_case							:	str			=						AutoImageProcessor.from_pretrained("""openmmlab/upernet-swin-tiny"""	)
   _snake_case							:	Any			=						UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-swin-tiny"""	).to(a_	)

   _snake_case							:	Dict			=						prepare_img()
   _snake_case							:	str			=						processor(images=a_,   return_tensors="""pt"""	).to(a_	)

   with torch.no_grad():
     _snake_case							:	Tuple			=						model(**a_	)

   _snake_case							:	Tuple			=						torch.Size((1, model.config.num_labels, 512, 512)	)
   self.assertEqual(outputs.logits.shape,   a_	)

   _snake_case							:	int			=						torch.tensor(
       [[-7.5_958, -7.5_958, -7.4_302], [-7.5_958, -7.5_958, -7.4_302], [-7.4_797, -7.4_797, -7.3_068]]	).to(a_	)
   self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3],   a_,   atol=1E-4	)	)


 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   _snake_case							:	Optional[Any]			=						AutoImageProcessor.from_pretrained("""openmmlab/upernet-convnext-tiny"""	)
   _snake_case							:	Optional[int]			=						UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-convnext-tiny"""	).to(a_	)

   _snake_case							:	List[str]			=						prepare_img()
   _snake_case							:	Tuple			=						processor(images=a_,   return_tensors="""pt"""	).to(a_	)

   with torch.no_grad():
     _snake_case							:	Optional[Any]			=						model(**a_	)

   _snake_case							:	Union[str, Any]			=						torch.Size((1, model.config.num_labels, 512, 512)	)
   self.assertEqual(outputs.logits.shape,   a_	)

   _snake_case							:	Optional[Any]			=						torch.tensor(
       [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]	).to(a_	)
   self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3],   a_,   atol=1E-4	)	)


