





import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD

def 				SCREAMING_SNAKE_CASE__							(	)   ->    List[str]:
  _lowercase						    =			'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
  _lowercase						    =			Image.open(requests.get(UpperCamelCase_      ,				stream=UpperCamelCase_  ).raw  ).convert('RGB'  )

  return image

def 				SCREAMING_SNAKE_CASE__							(	snake_case__	:Dict  )   ->    Any:
  _lowercase						    =			[]
  # fmt: off

  # vision encoder
  rename_keys.append(('visual_encoder.cls_token', 'vision_model.embeddings.class_embedding')  )
  rename_keys.append(('visual_encoder.pos_embed', 'vision_model.embeddings.position_embedding')  )
  rename_keys.append(('visual_encoder.patch_embed.proj.weight', 'vision_model.embeddings.patch_embedding.weight')  )
  rename_keys.append(('visual_encoder.patch_embed.proj.bias', 'vision_model.embeddings.patch_embedding.bias')  )
  rename_keys.append(('ln_vision.weight', 'vision_model.post_layernorm.weight')  )
  rename_keys.append(('ln_vision.bias', 'vision_model.post_layernorm.bias')  )

  for i in range(config.vision_config.num_hidden_layers  ):
    rename_keys.append((F"""visual_encoder.blocks.{i}.norm1.weight""", F"""vision_model.encoder.layers.{i}.layer_norm1.weight""")  )
    rename_keys.append((F"""visual_encoder.blocks.{i}.norm1.bias""", F"""vision_model.encoder.layers.{i}.layer_norm1.bias""")  )
    rename_keys.append((F"""visual_encoder.blocks.{i}.norm2.weight""", F"""vision_model.encoder.layers.{i}.layer_norm2.weight""")  )
    rename_keys.append((F"""visual_encoder.blocks.{i}.norm2.bias""", F"""vision_model.encoder.layers.{i}.layer_norm2.bias""")  )
    rename_keys.append((F"""visual_encoder.blocks.{i}.attn.qkv.weight""", F"""vision_model.encoder.layers.{i}.self_attn.qkv.weight""")  )
    rename_keys.append((F"""visual_encoder.blocks.{i}.attn.proj.weight""", F"""vision_model.encoder.layers.{i}.self_attn.projection.weight""",)  )
    rename_keys.append((F"""visual_encoder.blocks.{i}.attn.proj.bias""", F"""vision_model.encoder.layers.{i}.self_attn.projection.bias""")  )
    rename_keys.append((F"""visual_encoder.blocks.{i}.mlp.fc1.weight""", F"""vision_model.encoder.layers.{i}.mlp.fc1.weight""")  )
    rename_keys.append((F"""visual_encoder.blocks.{i}.mlp.fc1.bias""", F"""vision_model.encoder.layers.{i}.mlp.fc1.bias""")  )
    rename_keys.append((F"""visual_encoder.blocks.{i}.mlp.fc2.weight""", F"""vision_model.encoder.layers.{i}.mlp.fc2.weight""")  )
    rename_keys.append((F"""visual_encoder.blocks.{i}.mlp.fc2.bias""", F"""vision_model.encoder.layers.{i}.mlp.fc2.bias""")  )

  # QFormer
  rename_keys.append(('Qformer.bert.embeddings.LayerNorm.weight', 'qformer.layernorm.weight')  )
  rename_keys.append(('Qformer.bert.embeddings.LayerNorm.bias', 'qformer.layernorm.bias')  )

  # fmt: on
  return rename_keys

def 				SCREAMING_SNAKE_CASE__							(	snake_case__	:int      ,				snake_case__	:Optional[int]      ,				snake_case__	:int  )   ->    List[Any]:
  _lowercase						    =			dct.pop(UpperCamelCase_  )
  _lowercase						    =			val

def 				SCREAMING_SNAKE_CASE__							(	snake_case__	:Union[str, Any]      ,				snake_case__	:Optional[int]  )   ->    Optional[int]:
  for i in range(config.vision_config.num_hidden_layers  ):
    # read in original q and v biases
    _lowercase						    =			state_dict.pop(F"""visual_encoder.blocks.{i}.attn.q_bias"""  )
    _lowercase						    =			state_dict.pop(F"""visual_encoder.blocks.{i}.attn.v_bias"""  )

    # next, set bias in the state dict
    _lowercase						    =			torch.cat((q_bias, torch.zeros_like(UpperCamelCase_      ,				requires_grad=UpperCamelCase_  ), v_bias)  )
    _lowercase						    =			qkv_bias

def 				SCREAMING_SNAKE_CASE__							(	snake_case__	:List[Any]      ,				snake_case__	:List[Any]  )   ->    int:
  _lowercase						    =			364 if 'coco' in model_name else 224
  _lowercase						    =			BlipaVisionConfig(image_size=UpperCamelCase_  ).to_dict()

  # make sure the models have proper bos_token_id and eos_token_id set (important for generation)
  # seems like flan-T5 models don't have bos_token_id properly set?
  if "opt-2.7b" in model_name:
    _lowercase						    =			OPTConfig.from_pretrained('facebook/opt-2.7b'      ,				eos_token_id=UpperCamelCase_  ).to_dict()
  elif "opt-6.7b" in model_name:
    _lowercase						    =			OPTConfig.from_pretrained('facebook/opt-6.7b'      ,				eos_token_id=UpperCamelCase_  ).to_dict()
  elif "t5-xl" in model_name:
    _lowercase						    =			TaConfig.from_pretrained('google/flan-t5-xl'      ,				dense_act_fn='gelu'      ,				bos_token_id=1  ).to_dict()
  elif "t5-xxl" in model_name:
    _lowercase						    =			TaConfig.from_pretrained('google/flan-t5-xxl'      ,				dense_act_fn='gelu'      ,				bos_token_id=1  ).to_dict()

  _lowercase						    =			BlipaConfig(vision_config=UpperCamelCase_      ,				text_config=UpperCamelCase_  )

  return config, image_size

@torch.no_grad()
def 				SCREAMING_SNAKE_CASE__							(	snake_case__	:Tuple      ,				snake_case__	:Optional[Any]=None      ,				snake_case__	:Any=False  )   ->    Optional[Any]:
  _lowercase						    =			(
      AutoTokenizer.from_pretrained('facebook/opt-2.7b'  )
      if 'opt' in model_name
      else AutoTokenizer.from_pretrained('google/flan-t5-xl'  )
  )
  _lowercase						    =			tokenizer('\n'      ,				add_special_tokens=UpperCamelCase_  ).input_ids[0]
  _lowercase    ,		_lowercase						    =			get_blipa_config(UpperCamelCase_      ,				eos_token_id=UpperCamelCase_  )

  _lowercase						    =			BlipaForConditionalGeneration(UpperCamelCase_  ).eval()

  _lowercase						    =			{
      'blip2-opt-2.7b': ('blip2_opt', 'pretrain_opt2.7b'),
      'blip2-opt-6.7b': ('blip2_opt', 'pretrain_opt6.7b'),
      'blip2-opt-2.7b-coco': ('blip2_opt', 'caption_coco_opt2.7b'),
      'blip2-opt-6.7b-coco': ('blip2_opt', 'caption_coco_opt6.7b'),
      'blip2-flan-t5-xl': ('blip2_t5', 'pretrain_flant5xl'),
      'blip2-flan-t5-xl-coco': ('blip2_t5', 'caption_coco_flant5xl'),
      'blip2-flan-t5-xxl': ('blip2_t5', 'pretrain_flant5xxl'),
  }

  _lowercase    ,		_lowercase						    =			model_name_to_original[model_name]

  # load original model
  print('Loading original model...'  )
  _lowercase						    =			'cuda' if torch.cuda.is_available() else 'cpu'
  _lowercase    ,		_lowercase    ,		_lowercase						    =			load_model_and_preprocess(
      name=UpperCamelCase_      ,				model_type=UpperCamelCase_      ,				is_eval=UpperCamelCase_      ,				device=UpperCamelCase_  )
  original_model.eval()
  print('Done!'  )

  # update state dict keys
  _lowercase						    =			original_model.state_dict()
  _lowercase						    =			create_rename_keys(UpperCamelCase_  )
  for src, dest in rename_keys:
    rename_key(UpperCamelCase_      ,				UpperCamelCase_      ,				UpperCamelCase_  )

  # some keys can be renamed efficiently
  for key, val in state_dict.copy().items():
    _lowercase						    =			state_dict.pop(UpperCamelCase_  )
    if key.startswith('Qformer.bert'  ):
      _lowercase						    =			key.replace('Qformer.bert'      ,				'qformer'  )
    if "attention.self" in key:
      _lowercase						    =			key.replace('self'      ,				'attention'  )
    if "opt_proj" in key:
      _lowercase						    =			key.replace('opt_proj'      ,				'language_projection'  )
    if "t5_proj" in key:
      _lowercase						    =			key.replace('t5_proj'      ,				'language_projection'  )
    if key.startswith('opt'  ):
      _lowercase						    =			key.replace('opt'      ,				'language'  )
    if key.startswith('t5'  ):
      _lowercase						    =			key.replace('t5'      ,				'language'  )
    _lowercase						    =			val

  # read in qv biases
  read_in_q_v_bias(UpperCamelCase_      ,				UpperCamelCase_  )

  _lowercase    ,		_lowercase						    =			hf_model.load_state_dict(UpperCamelCase_      ,				strict=UpperCamelCase_  )
  assert len(UpperCamelCase_  ) == 0
  assert unexpected_keys == ["qformer.embeddings.position_ids"]

  _lowercase						    =			load_demo_image()
  _lowercase						    =			vis_processors['eval'](UpperCamelCase_  ).unsqueeze(0  ).to(UpperCamelCase_  )
  _lowercase						    =			tokenizer(['\n']      ,				return_tensors='pt'  ).input_ids.to(UpperCamelCase_  )

  # create processor
  _lowercase						    =			BlipImageProcessor(
      size={'height': image_size, 'width': image_size}      ,				image_mean=UpperCamelCase_      ,				image_std=UpperCamelCase_  )
  _lowercase						    =			BlipaProcessor(image_processor=UpperCamelCase_      ,				tokenizer=UpperCamelCase_  )
  _lowercase						    =			processor(images=UpperCamelCase_      ,				return_tensors='pt'  ).pixel_values.to(UpperCamelCase_  )

  # make sure processor creates exact same pixel values
  assert torch.allclose(UpperCamelCase_      ,				UpperCamelCase_  )

  original_model.to(UpperCamelCase_  )
  hf_model.to(UpperCamelCase_  )
  with torch.no_grad():
    if "opt" in model_name:
      _lowercase						    =			original_model({'image': original_pixel_values, 'text_input': ['']}  ).logits
      _lowercase						    =			hf_model(UpperCamelCase_      ,				UpperCamelCase_  ).logits
    else:
      _lowercase						    =			original_model(
          {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}  ).logits
      _lowercase						    =			input_ids.masked_fill(input_ids == tokenizer.pad_token_id      ,				-100  )
      _lowercase						    =			hf_model(UpperCamelCase_      ,				UpperCamelCase_      ,				labels=UpperCamelCase_  ).logits

  assert original_logits.shape == logits.shape
  print('First values of original logits:'      ,				original_logits[0, :3, :3]  )
  print('First values of HF logits:'      ,				logits[0, :3, :3]  )

  # assert values
  if model_name == "blip2-flan-t5-xl":
    _lowercase						    =			torch.tensor(
        [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]      ,				device=UpperCamelCase_  )
    assert torch.allclose(logits[0, :3, :3]      ,				UpperCamelCase_      ,				atol=1E-4  )
  elif model_name == "blip2-flan-t5-xl-coco":
    _lowercase						    =			torch.tensor(
        [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]      ,				device=UpperCamelCase_  )
  else:
    # cast to same type
    _lowercase						    =			logits.dtype
    assert torch.allclose(original_logits.to(UpperCamelCase_  )      ,				UpperCamelCase_      ,				atol=1E-2  )
  print('Looks ok!'  )

  print('Generating a caption...'  )
  _lowercase						    =			''
  _lowercase						    =			tokenizer(UpperCamelCase_      ,				return_tensors='pt'  ).input_ids.to(UpperCamelCase_  )

  _lowercase						    =			original_model.generate({'image': original_pixel_values}  )
  _lowercase						    =			hf_model.generate(
      UpperCamelCase_      ,				UpperCamelCase_      ,				do_sample=UpperCamelCase_      ,				num_beams=5      ,				max_length=30      ,				min_length=1      ,				top_p=0.9      ,				repetition_penalty=1.0      ,				length_penalty=1.0      ,				temperature=1      ,				)
  print('Original generation:'      ,				UpperCamelCase_  )
  _lowercase						    =			input_ids.shape[1]
  _lowercase						    =			processor.batch_decode(outputs[:, prompt_length:]      ,				skip_special_tokens=UpperCamelCase_  )
  _lowercase						    =			[text.strip() for text in output_text]
  print('HF generation:'      ,				UpperCamelCase_  )

  if pytorch_dump_folder_path is not None:
    processor.save_pretrained(UpperCamelCase_  )
    hf_model.save_pretrained(UpperCamelCase_  )

  if push_to_hub:
    processor.push_to_hub(F"""nielsr/{model_name}"""  )
    hf_model.push_to_hub(F"""nielsr/{model_name}"""  )


if __name__ == "__main__":
      snake_case								=							argparse.ArgumentParser()
      snake_case								=							[
          """blip2-opt-2.7b""",
          """blip2-opt-6.7b""",
          """blip2-opt-2.7b-coco""",
          """blip2-opt-6.7b-coco""",
          """blip2-flan-t5-xl""",
          """blip2-flan-t5-xl-coco""",
          """blip2-flan-t5-xxl""",
      ]
      parser.add_argument(
          """--model_name""",
          default="""blip2-opt-2.7b""",
          choices=choices,
          type=str,
          help="""Path to hf config.json of model to convert""",
      )
      parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
      parser.add_argument(
          """--push_to_hub""",
          action="""store_true""",
          help="""Whether to push the model and processor to the hub after converting""",
      )

      snake_case								=							parser.parse_args()

      convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




from __future__ import annotations

import unittest

from transformers import DebertaVaConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
   import tensorflow as tf

   from transformers import (
       TFDebertaVaForMaskedLM,
       TFDebertaVaForQuestionAnswering,
       TFDebertaVaForSequenceClassification,
       TFDebertaVaForTokenClassification,
       TFDebertaVaModel,
   )



class      _UpperCAmelCase      :




     def __init__(					self			,    a__			,    a__=1_3			,    a__=7			,    a__=True			,    a__=True			,    a__=True			,    a__=True			,    a__=9_9			,    a__=3_2			,    a__=2			,    a__=4			,    a__=3_7			,    a__="gelu"			,    a__=0.1			,    a__=0.1			,    a__=5_1_2			,    a__=1_6			,    a__=2			,    a__=0.0_2			,    a__=False			,    a__=True			,    a__="None"			,    a__=3			,    a__=4			,    a__=None			,    ):
         A__					=       parent
         A__					=       batch_size
         A__					=       seq_length
         A__					=       is_training
         A__					=       use_input_mask
         A__					=       use_token_type_ids
         A__					=       use_labels
         A__					=       vocab_size
         A__					=       hidden_size
         A__					=       num_hidden_layers
         A__					=       num_attention_heads
         A__					=       intermediate_size
         A__					=       hidden_act
         A__					=       hidden_dropout_prob
         A__					=       attention_probs_dropout_prob
         A__					=       max_position_embeddings
         A__					=       type_vocab_size
         A__					=       type_sequence_label_size
         A__					=       initializer_range
         A__					=       num_labels
         A__					=       num_choices
         A__					=       relative_attention
         A__					=       position_biased_input
         A__					=       pos_att_type
         A__					=       scope




     def 		snake_case_		(					self):
         A__					=       ids_tensor([self.batch_size, self.seq_length]			,    self.vocab_size)

         A__					=       None
         if self.use_input_mask:
             A__					=       random_attention_mask([self.batch_size, self.seq_length])

         A__					=       None
         if self.use_token_type_ids:
             A__					=       ids_tensor([self.batch_size, self.seq_length]			,    self.type_vocab_size)

         A__					=       None
         A__					=       None
         A__					=       None
         if self.use_labels:
             A__					=       ids_tensor([self.batch_size]			,    self.type_sequence_label_size)
             A__					=       ids_tensor([self.batch_size, self.seq_length]			,    self.num_labels)

         A__					=       DebertaVaConfig(
             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			,    relative_attention=self.relative_attention			,    position_biased_input=self.position_biased_input			,    initializer_range=self.initializer_range			,    return_dict=a__			,    )

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




     def 		snake_case_		(					self			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__):
         A__					=       TFDebertaVaModel(config=a__)
         A__					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}

         A__					=       [input_ids, input_mask]
         A__					=       model(a__)

         A__					=       model(a__)

         self.parent.assertEqual(result.last_hidden_state.shape			,    (self.batch_size, self.seq_length, self.hidden_size))




     def 		snake_case_		(					self			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__):
         A__					=       TFDebertaVaForMaskedLM(config=a__)
         A__					=       {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }
         A__					=       model(a__)
         self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.seq_length, self.vocab_size))




     def 		snake_case_		(					self			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__):
         A__					=       self.num_labels
         A__					=       TFDebertaVaForSequenceClassification(config=a__)
         A__					=       {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }

         A__					=       model(a__)
         self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.num_labels))




     def 		snake_case_		(					self			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__):
         A__					=       self.num_labels
         A__					=       TFDebertaVaForTokenClassification(config=a__)
         A__					=       {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }
         A__					=       model(a__)
         self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.seq_length, self.num_labels))




     def 		snake_case_		(					self			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__):
         A__					=       TFDebertaVaForQuestionAnswering(config=a__)
         A__					=       {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }

         A__					=       model(a__)
         self.parent.assertEqual(result.start_logits.shape			,    (self.batch_size, self.seq_length))
         self.parent.assertEqual(result.end_logits.shape			,    (self.batch_size, self.seq_length))




     def 		snake_case_		(					self):
         A__					=       self.prepare_config_and_inputs()
         (
             (
             A__
         )			,						(
             A__
         )			,						(
             A__
         )			,						(
             A__
         )			,						(
             A__
         )			,						(
             A__
         )			,						(
             A__
         )			,						
         )					=       config_and_inputs
         A__					=       {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
         return config, inputs_dict







@require_tf
class      _UpperCAmelCase      (							A__      ,   A__      ,   unittest.TestCase  ):
     UpperCamelCase__		     =					(
         (
             TFDebertaVaModel,
             TFDebertaVaForMaskedLM,
             TFDebertaVaForQuestionAnswering,
             TFDebertaVaForSequenceClassification,
             TFDebertaVaForTokenClassification,
         )
         if is_tf_available()
         else ()
     )
     UpperCamelCase__		     =					(
         {
             '''feature-extraction''': TFDebertaVaModel,
             '''fill-mask''': TFDebertaVaForMaskedLM,
             '''question-answering''': TFDebertaVaForQuestionAnswering,
             '''text-classification''': TFDebertaVaForSequenceClassification,
             '''token-classification''': TFDebertaVaForTokenClassification,
             '''zero-shot''': TFDebertaVaForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )

     UpperCamelCase__		     =					False
     UpperCamelCase__		     =					False




     def 		snake_case_		(					self):
         A__					=       TFDebertaVaModelTester(self)
         A__					=       ConfigTester(self			,    config_class=a__			,    hidden_size=3_7)




     def 		snake_case_		(					self):
         self.config_tester.run_common_tests()




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a__)




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*a__)




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*a__)




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*a__)




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*a__)




     @slow
     def 		snake_case_		(					self):
         A__					=       TFDebertaVaModel.from_pretrained('''kamalkraj/deberta-v2-xlarge''')
         self.assertIsNotNone(a__)







@require_tf
class      _UpperCAmelCase      (							unittest.TestCase  ):




     @unittest.skip(reason='''Model not available yet''')
     def 		snake_case_		(					self):
         pass




     @slow
     def 		snake_case_		(					self):
         A__					=       TFDebertaVaModel.from_pretrained('''kamalkraj/deberta-v2-xlarge''')
         A__					=       tf.constant([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]])
         A__					=       tf.constant([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
         A__					=       model(a__			,    attention_mask=a__)[0]

         A__					=       tf.constant(
             [[[0.2_3_5_6, 0.1_9_4_8, 0.0_3_6_9], [-0.1_0_6_3, 0.3_5_8_6, -0.5_1_5_2], [-0.6_3_9_9, -0.0_2_5_9, -0.2_5_2_5]]])
         tf.debugging.assert_near(output[:, 1:4, 1:4]			,    a__			,    atol=1e-4)
