





from __future__ import annotations

import unittest

from transformers import BlenderbotConfig, BlenderbotTokenizer, is_tf_available
from transformers.testing_utils import require_tf, require_tokenizers, slow
from transformers.utils import cached_property

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


if is_tf_available():
    import tensorflow as tf

    from transformers import TFAutoModelForSeqaSeqLM, TFBlenderbotForConditionalGeneration, TFBlenderbotModel



@require_tf
class    A__  :
      _UpperCAmelCase      :Any		       =  BlenderbotConfig
      _UpperCAmelCase      :str		       =  {}
      _UpperCAmelCase      :List[str]		       =  'gelu'





      def __init__(					self			,				A_			,				A_=13			,				A_=7			,				A_=True			,				A_=False			,				A_=99			,				A_=32			,				A_=2			,				A_=4			,				A_=37			,				A_=0.1			,				A_=0.1			,				A_=20			,				A_=2			,				A_=1			,				A_=0			,				):
       '''simple docstring'''




       UpperCamelCase      :   Dict			      =  parent
       UpperCamelCase      :   Union[str, Any]			      =  batch_size
       UpperCamelCase      :   str			      =  seq_length
       UpperCamelCase      :   int			      =  is_training
       UpperCamelCase      :   List[str]			      =  use_labels
       UpperCamelCase      :   Optional[Any]			      =  vocab_size
       UpperCamelCase      :   str			      =  hidden_size
       UpperCamelCase      :   List[Any]			      =  num_hidden_layers
       UpperCamelCase      :   Optional[int]			      =  num_attention_heads
       UpperCamelCase      :   Tuple			      =  intermediate_size

       UpperCamelCase      :   Optional[int]			      =  hidden_dropout_prob
       UpperCamelCase      :   Union[str, Any]			      =  attention_probs_dropout_prob
       UpperCamelCase      :   List[Any]			      =  max_position_embeddings
       UpperCamelCase      :   str			      =  eos_token_id
       UpperCamelCase      :   Dict			      =  pad_token_id
       UpperCamelCase      :   Tuple			      =  bos_token_id





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   List[str]			      =  ids_tensor([self.batch_size, self.seq_length - 1]			,				self.vocab_size  )
       UpperCamelCase      :   Any			      =  tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size  )			,				1  )
       UpperCamelCase      :   Union[str, Any]			      =  tf.concat([input_ids, eos_tensor]			,				axis=1  )

       UpperCamelCase      :   Union[str, Any]			      =  ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size  )

       UpperCamelCase      :   List[Any]			      =  self.config_cls(
           vocab_size=self.vocab_size			,				d_model=self.hidden_size			,				encoder_layers=self.num_hidden_layers			,				decoder_layers=self.num_hidden_layers			,				encoder_attention_heads=self.num_attention_heads			,				decoder_attention_heads=self.num_attention_heads			,				encoder_ffn_dim=self.intermediate_size			,				decoder_ffn_dim=self.intermediate_size			,				dropout=self.hidden_dropout_prob			,				attention_dropout=self.attention_probs_dropout_prob			,				max_position_embeddings=self.max_position_embeddings			,				eos_token_ids=[2]			,				bos_token_id=self.bos_token_id			,				pad_token_id=self.pad_token_id			,				decoder_start_token_id=self.pad_token_id			,				**self.config_updates			,				)
       UpperCamelCase      :   Optional[int]			      =  prepare_blenderbot_inputs_dict(A_			,				A_			,				A_  )
       return config, inputs_dict





      def 			__UpperCamelCase(					self			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   List[str]			      =  TFBlenderbotModel(config=A_  ).get_decoder()
       UpperCamelCase      :   Optional[Any]			      =  inputs_dict["input_ids"]

       UpperCamelCase      :   int			      =  input_ids[:1, :]
       UpperCamelCase      :   Any			      =  inputs_dict["attention_mask"][:1, :]
       UpperCamelCase      :   Tuple			      =  inputs_dict["head_mask"]
       UpperCamelCase      :   str			      =  1

       # first forward pass
       UpperCamelCase      :   List[Any]			      =  model(A_			,				attention_mask=A_			,				head_mask=A_			,				use_cache=A_  )

       UpperCamelCase      ,       UpperCamelCase      :   List[Any]			      =  outputs.to_tuple()

       # create hypothetical next token and extent to next_input_ids
       UpperCamelCase      :   List[str]			      =  ids_tensor((self.batch_size, 3)			,				config.vocab_size  )
       UpperCamelCase      :   List[Any]			      =  tf.cast(ids_tensor((self.batch_size, 3)			,				2  )			,				tf.inta  )

       # append to next input_ids and
       UpperCamelCase      :   Union[str, Any]			      =  tf.concat([input_ids, next_tokens]			,				axis=-1  )
       UpperCamelCase      :   Union[str, Any]			      =  tf.concat([attention_mask, next_attn_mask]			,				axis=-1  )

       UpperCamelCase      :   List[str]			      =  model(A_			,				attention_mask=A_  )[0]
       UpperCamelCase      :   List[Any]			      =  model(A_			,				attention_mask=A_			,				past_key_values=A_  )[0]

       self.parent.assertEqual(next_tokens.shape[1]			,				output_from_past.shape[1]  )

       # select random slice
       UpperCamelCase      :   List[Any]			      =  int(ids_tensor((1,)			,				output_from_past.shape[-1]  )  )
       UpperCamelCase      :   List[str]			      =  output_from_no_past[:, -3:, random_slice_idx]
       UpperCamelCase      :   Dict			      =  output_from_past[:, :, random_slice_idx]

       # test that outputs are equal for slice
       tf.debugging.assert_near(A_			,				A_			,				rtol=1e-3  )



def 						A_							(	_lowerCAmelCase				,  _lowerCAmelCase				,  _lowerCAmelCase				,  _lowerCAmelCase=None				,  _lowerCAmelCase=None				,  _lowerCAmelCase=None				,  _lowerCAmelCase=None				,  _lowerCAmelCase=None				,  )		->					Tuple:
 if attention_mask is None:
  UpperCamelCase      :   Union[str, Any]			      =  tf.cast(tf.math.not_equal(_lowerCAmelCase				,  config.pad_token_id	)				,  tf.inta	)
 if decoder_attention_mask is None:
  UpperCamelCase      :   List[str]			      =  tf.concat(
      [
          tf.ones(decoder_input_ids[:, :1].shape				,  dtype=tf.inta	),
          tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]				,  config.pad_token_id	)				,  tf.inta	),
      ]				,  axis=-1				,  )
 if head_mask is None:
  UpperCamelCase      :   List[str]			      =  tf.ones((config.encoder_layers, config.encoder_attention_heads)	)
 if decoder_head_mask is None:
  UpperCamelCase      :   Union[str, Any]			      =  tf.ones((config.decoder_layers, config.decoder_attention_heads)	)
 if cross_attn_head_mask is None:
  UpperCamelCase      :   Any			      =  tf.ones((config.decoder_layers, config.decoder_attention_heads)	)
 return {
     "input_ids": input_ids,
     "decoder_input_ids": decoder_input_ids,
     "attention_mask": attention_mask,
     "decoder_attention_mask": decoder_attention_mask,
     "head_mask": head_mask,
     "decoder_head_mask": decoder_head_mask,
     "cross_attn_head_mask": cross_attn_head_mask,
 }



@require_tf
class    A__  (			__snake_case			,  __snake_case			,  unittest.TestCase     ):
      _UpperCAmelCase      :str		       =  (TFBlenderbotForConditionalGeneration, TFBlenderbotModel) if is_tf_available() else ()
      _UpperCAmelCase      :List[str]		       =  (TFBlenderbotForConditionalGeneration,) if is_tf_available() else ()
      _UpperCAmelCase      :Optional[Any]		       =  (
          {
              'conversational': TFBlenderbotForConditionalGeneration,
              'feature-extraction': TFBlenderbotModel,
              'summarization': TFBlenderbotForConditionalGeneration,
              'text2text-generation': TFBlenderbotForConditionalGeneration,
              'translation': TFBlenderbotForConditionalGeneration,
          }
          if is_tf_available()
          else {}
      )
      _UpperCAmelCase      :Optional[Any]		       =  True
      _UpperCAmelCase      :Any		       =  False
      _UpperCAmelCase      :List[str]		       =  False





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Tuple			      =  TFBlenderbotModelTester(self  )
       UpperCamelCase      :   Any			      =  ConfigTester(self			,				config_class=A_  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Optional[Any]			      =  self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.check_decoder_model_past_large_inputs(*A_  )



@require_tokenizers
@require_tf
class    A__  (			unittest.TestCase     ):
      _UpperCAmelCase      :Optional[Any]		       =  ['My friends are cool but they eat too many carbs.']
      _UpperCAmelCase      :Tuple		       =  'facebook/blenderbot-400M-distill'





      @cached_property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       return BlenderbotTokenizer.from_pretrained(self.model_name  )





      @cached_property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Optional[int]			      =  TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name  )
       return model





      @slow
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   str			      =  self.tokenizer(self.src_text			,				return_tensors="tf"  )
       UpperCamelCase      :   Optional[int]			      =  self.model.generate(
           model_inputs.input_ids			,				)
       UpperCamelCase      :   Union[str, Any]			      =  self.tokenizer.batch_decode(generated_ids.numpy()			,				skip_special_tokens=A_  )[0]
       assert (
           generated_words
           == " That's unfortunate. Are they trying to lose weight or are they just trying to be healthier?"
       )







import inspect
import math
import tempfile
import unittest

import numpy as np

from transformers import ViTMAEConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
    import torch
    from torch import nn

    from transformers import ViTMAEForPreTraining, ViTMAEModel
    from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import ViTImageProcessor



class    A__  :





      def __init__(					self			,				A_			,				A_=13			,				A_=30			,				A_=2			,				A_=3			,				A_=True			,				A_=True			,				A_=32			,				A_=5			,				A_=4			,				A_=37			,				A_="gelu"			,				A_=0.1			,				A_=0.1			,				A_=10			,				A_=0.02			,				A_=3			,				A_=0.6			,				A_=None			,				):
       '''simple docstring'''




       UpperCamelCase      :   Dict			      =  parent
       UpperCamelCase      :   Union[str, Any]			      =  batch_size
       UpperCamelCase      :   Any			      =  image_size
       UpperCamelCase      :   Tuple			      =  patch_size
       UpperCamelCase      :   Union[str, Any]			      =  num_channels
       UpperCamelCase      :   Tuple			      =  is_training
       UpperCamelCase      :   int			      =  use_labels
       UpperCamelCase      :   Tuple			      =  hidden_size
       UpperCamelCase      :   Optional[Any]			      =  num_hidden_layers
       UpperCamelCase      :   str			      =  num_attention_heads
       UpperCamelCase      :   Optional[int]			      =  intermediate_size
       UpperCamelCase      :   Dict			      =  hidden_act
       UpperCamelCase      :   List[str]			      =  hidden_dropout_prob
       UpperCamelCase      :   str			      =  attention_probs_dropout_prob
       UpperCamelCase      :   Optional[Any]			      =  type_sequence_label_size
       UpperCamelCase      :   Tuple			      =  initializer_range
       UpperCamelCase      :   Any			      =  mask_ratio
       UpperCamelCase      :   int			      =  scope

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





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




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

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

       UpperCamelCase      :   str			      =  self.get_config()

       return config, pixel_values, labels





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       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			,				hidden_act=self.hidden_act			,				hidden_dropout_prob=self.hidden_dropout_prob			,				attention_probs_dropout_prob=self.attention_probs_dropout_prob			,				is_decoder=A_			,				initializer_range=self.initializer_range			,				mask_ratio=self.mask_ratio			,				)





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   Any			      =  ViTMAEModel(config=A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   Optional[Any]			      =  model(A_  )
       self.parent.assertEqual(result.last_hidden_state.shape			,				(self.batch_size, self.seq_length, self.hidden_size)  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       UpperCamelCase      :   Tuple			      =  ViTMAEForPreTraining(A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   List[str]			      =  model(A_  )
       UpperCamelCase      :   Any			      =  (self.image_size // self.patch_size) ** 2
       UpperCamelCase      :   List[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
       UpperCamelCase      :   str			      =  1
       UpperCamelCase      :   Union[str, Any]			      =  ViTMAEForPreTraining(A_  )
       model.to(A_  )
       model.eval()
       UpperCamelCase      :   Union[str, Any]			      =  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
       UpperCamelCase      :   Union[str, Any]			      =  model(A_  )
       UpperCamelCase      :   List[str]			      =  self.patch_size**2
       self.parent.assertEqual(result.logits.shape			,				(self.batch_size, num_patches, expected_num_channels)  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Union[str, Any]			      =  self.prepare_config_and_inputs()
       UpperCamelCase      ,       UpperCamelCase      ,       UpperCamelCase      :   Optional[int]			      =  config_and_inputs
       UpperCamelCase      :   Optional[int]			      =  {"pixel_values": pixel_values}
       return config, inputs_dict



@require_torch
class    A__  (			__snake_case			,  __snake_case			,  unittest.TestCase     ):
      _UpperCAmelCase      :Any		       =  (ViTMAEModel, ViTMAEForPreTraining) if is_torch_available() else ()
      _UpperCAmelCase      :int		       =  {'feature-extraction': ViTMAEModel} if is_torch_available() else {}

      _UpperCAmelCase      :Optional[Any]		       =  False
      _UpperCAmelCase      :Any		       =  False
      _UpperCAmelCase      :Optional[int]		       =  False
      _UpperCAmelCase      :Dict		       =  False





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Union[str, Any]			      =  ViTMAEModelTester(self  )
       UpperCamelCase      :   Any			      =  ConfigTester(self			,				config_class=A_			,				has_text_modality=A_			,				hidden_size=37  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      @unittest.skip(reason="ViTMAE does not use inputs_embeds"  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      ,       UpperCamelCase      :   List[Any]			      =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCamelCase      :   Any			      =  model_class(A_  )
        self.assertIsInstance(model.get_input_embeddings()			,				(nn.Module)  )
        UpperCamelCase      :   Any			      =  model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(A_			,				nn.Linear  )  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      ,       UpperCamelCase      :   Optional[int]			      =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCamelCase      :   List[str]			      =  model_class(A_  )
        UpperCamelCase      :   Any			      =  inspect.signature(model.forward  )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        UpperCamelCase      :   List[str]			      =  [*signature.parameters.keys()]

        UpperCamelCase      :   Optional[int]			      =  ["pixel_values"]
        self.assertListEqual(arg_names[:1]			,				A_  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   List[str]			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Optional[int]			      =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_pretraining(*A_  )





      def 			__UpperCamelCase(					self			,				A_			,				A_			,				A_  ):
       '''simple docstring'''




       np.random.seed(2  )

       UpperCamelCase      :   Union[str, Any]			      =  int((pt_model.config.image_size // pt_model.config.patch_size) ** 2  )
       UpperCamelCase      :   Optional[Any]			      =  np.random.uniform(size=(self.model_tester.batch_size, num_patches)  )
       UpperCamelCase      :   Tuple			      =  torch.from_numpy(A_  )

       # Add `noise` argument.
       # PT inputs will be prepared in `super().check_pt_tf_models()` with this added `noise` argument
       UpperCamelCase      :   Any			      =  pt_noise

       super().check_pt_tf_models(A_			,				A_			,				A_  )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      ,       UpperCamelCase      :   Tuple			      =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCamelCase      :   Optional[Any]			      =  model_class(A_  )
        model.to(A_  )
        model.eval()
        # make random mask reproducible
        torch.manual_seed(2  )
        with torch.no_grad():
         UpperCamelCase      :   str			      =  model(**self._prepare_for_class(A_			,				A_  )  )

        UpperCamelCase      :   Any			      =  outputs[0].cpu().numpy()
        UpperCamelCase      :   str			      =  0

        with tempfile.TemporaryDirectory() as tmpdirname:
         model.save_pretrained(A_  )
         UpperCamelCase      :   Optional[int]			      =  model_class.from_pretrained(A_  )
         model.to(A_  )
         # make random mask reproducible
         torch.manual_seed(2  )
         with torch.no_grad():
          UpperCamelCase      :   Tuple			      =  model(**self._prepare_for_class(A_			,				A_  )  )

         # Make sure we don't have nans
         UpperCamelCase      :   Dict			      =  after_outputs[0].cpu().numpy()
         UpperCamelCase      :   int			      =  0
         UpperCamelCase      :   str			      =  np.amax(np.abs(out_a - out_a  )  )
         self.assertLessEqual(A_			,				1e-5  )





      @unittest.skip(
          reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass





      @unittest.skip(
          reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass





      @unittest.skip(
          reason="ViTMAE returns a random mask + ids_restore in each forward pass. See test_save_load\n    to get deterministic results."  )
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       pass





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




       pass





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




       pass





      @slow
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCamelCase      :   List[str]			      =  ViTMAEModel.from_pretrained(A_  )
        self.assertIsNotNone(A_  )



def 						A_							(	)		->					Optional[int]:
 UpperCamelCase      :   List[Any]			      =  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"	)
 return image



@require_torch
@require_vision
class    A__  (			unittest.TestCase     ):





      @cached_property
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       return ViTImageProcessor.from_pretrained("facebook/vit-mae-base"  ) if is_vision_available() else None





      @slow
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       np.random.seed(2  )

       UpperCamelCase      :   Union[str, Any]			      =  ViTMAEForPreTraining.from_pretrained("facebook/vit-mae-base"  ).to(A_  )

       UpperCamelCase      :   Any			      =  self.default_image_processor
       UpperCamelCase      :   Optional[Any]			      =  prepare_img()
       UpperCamelCase      :   int			      =  image_processor(images=A_			,				return_tensors="pt"  ).to(A_  )

       # 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)
       UpperCamelCase      :   Dict			      =  ViTMAEConfig()
       UpperCamelCase      :   Any			      =  int((vit_mae_config.image_size // vit_mae_config.patch_size) ** 2  )
       UpperCamelCase      :   Optional[Any]			      =  np.random.uniform(size=(1, num_patches)  )

       # forward pass
       with torch.no_grad():
        UpperCamelCase      :   Optional[Any]			      =  model(**A_			,				noise=torch.from_numpy(A_  ).to(device=A_  )  )

       # verify the logits
       UpperCamelCase      :   str			      =  torch.Size((1, 196, 768)  )
       self.assertEqual(outputs.logits.shape			,				A_  )

       UpperCamelCase      :   str			      =  torch.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]]  )

       self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]			,				expected_slice.to(A_  )			,				atol=1e-4  )  )
