




'''simple docstring'''





import json
import os
import subprocess
import sys
import warnings
from argparse import ArgumentParser
from contextlib import AbstractContextManager
from typing import Dict, List, Optional

import requests

from ..utils import logging
from . import BaseTransformersCLICommand


lowercase__								=						logging.get_logger(__name__)  # pylint: disable=invalid-name


lowercase__								=						'''lfs-multipart-upload'''






class     lowercase  ( UpperCamelCase_):






      @staticmethod
      def  A__(a__       )							-> List[str]:
         SCREAMING_SNAKE_CASE__      :					List[Any]							    =			parser.add_parser(
             '''lfs-enable-largefiles'''			,	help=(
                 '''Deprecated: use `huggingface-cli` instead. Configure your repository to enable upload of files > 5GB.'''
             )			,	)
         enable_parser.add_argument('''path'''			,	type=__a			,	help='''Local path to repository you want to configure.'''       )
         enable_parser.set_defaults(func=lambda a__       : LfsEnableCommand(__a       )       )

         SCREAMING_SNAKE_CASE__      :					Optional[int]							    =			parser.add_parser(
             __a			,	help=(
                 '''Deprecated: use `huggingface-cli` instead. '''
                 '''Command will get called by git-lfs, do not call it directly.'''
             )			,	)
         upload_parser.set_defaults(func=lambda a__       : LfsUploadCommand(__a       )       )






class     lowercase  :






      def __init__(self			,	a__       )							-> Union[str, Any]:
         SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			args






      def  A__(self       )							-> Tuple:
         warnings.warn(
             '''Managing repositories through transformers-cli is deprecated. Please use `huggingface-cli` instead.'''       )
         SCREAMING_SNAKE_CASE__      :					str							    =			os.path.abspath(self.args.path       )
         if not os.path.isdir(__a       ):
            print('''This does not look like a valid git repo.'''       )
            exit(1       )
         subprocess.run(
             '''git config lfs.customtransfer.multipart.path transformers-cli'''.split()			,	check=__a			,	cwd=__a       )
         subprocess.run(
             f"""git config lfs.customtransfer.multipart.args {LFS_MULTIPART_UPLOAD_COMMAND}""".split()			,	check=__a			,	cwd=__a			,	)
         print('''Local repo set up for largefiles'''       )


def 					UpperCamelCase						(	snake_case ):




   '''simple docstring'''

   SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			json.dumps(snake_case ) + '\n'
   sys.stdout.write(snake_case )
   sys.stdout.flush()


def 					UpperCamelCase						(	):




   '''simple docstring'''

   SCREAMING_SNAKE_CASE__      :					int							    =			json.loads(sys.stdin.readline().strip() )

   if "terminate" in (msg.get('''type''' ), msg.get('''event''' )):
      # terminate message received
      return None

   if msg.get('''event''' ) not in ("download", "upload"):
      logger.critical('''Received unexpected message''' )
      sys.exit(1 )

   return msg






class     lowercase  ( UpperCamelCase_):






      def __init__(self			,	a__			,	a__			,	a__       )							-> Dict:
         SCREAMING_SNAKE_CASE__      :					int							    =			filepath
         SCREAMING_SNAKE_CASE__      :					Dict							    =			seek_from
         SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			read_limit
         SCREAMING_SNAKE_CASE__      :					List[Any]							    =			0






      def __enter__(self       )							-> Optional[Any]:
         SCREAMING_SNAKE_CASE__      :					List[str]							    =			open(self.filepath			,	'''rb'''       )
         self.f.seek(self.seek_from       )
         return self






      def __len__(self       )							-> Any:
         SCREAMING_SNAKE_CASE__      :					Dict							    =			os.fstat(self.f.fileno()       ).st_size
         return min(self.read_limit			,	total_length - self.seek_from       )






      def  A__(self			,	a__=-1       )							-> Optional[Any]:






         if self.n_seen >= self.read_limit:
            return b""
         SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			self.read_limit - self.n_seen
         SCREAMING_SNAKE_CASE__      :					Tuple							    =			self.f.read(remaining_amount if n < 0 else min(__a			,	__a       )       )
         self.n_seen += len(__a       )
         return data






      def __iter__(self       )							-> List[Any]:
         yield self.read(n=4 * 1_0_2_4 * 1_0_2_4       )






      def __exit__(self			,	*a__       )							-> Tuple:
         self.f.close()






class     lowercase  :






      def __init__(self			,	a__       )							-> Any:
         SCREAMING_SNAKE_CASE__      :					List[Any]							    =			args






      def  A__(self       )							-> List[Any]:
         SCREAMING_SNAKE_CASE__      :					Dict							    =			json.loads(sys.stdin.readline().strip()       )
         if not (init_msg.get('''event'''       ) == "init" and init_msg.get('''operation'''       ) == "upload"):
            write_msg({'''error''': {'''code''': 3_2, '''message''': '''Wrong lfs init operation'''}}       )
            sys.exit(1       )

         # The transfer process should use the information it needs from the
         # initiation structure, and also perform any one-off setup tasks it
         # needs to do. It should then respond on stdout with a simple empty
         # confirmation structure, as follows:
         write_msg({}       )

         # After the initiation exchange, git-lfs will send any number of
         # transfer requests to the stdin of the transfer process, in a serial sequence.
         while True:
            SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			read_msg()
            if msg is None:
               # When all transfers have been processed, git-lfs will send
               # a terminate event to the stdin of the transfer process.
               # On receiving this message the transfer process should
               # clean up and terminate. No response is expected.
               sys.exit(0       )

            SCREAMING_SNAKE_CASE__      :					List[str]							    =			msg['oid']
            SCREAMING_SNAKE_CASE__      :					Union[str, Any]							    =			msg['path']
            SCREAMING_SNAKE_CASE__      :					List[Any]							    =			msg['action']['href']
            SCREAMING_SNAKE_CASE__      :					int							    =			msg['action']['header']
            SCREAMING_SNAKE_CASE__      :					Dict							    =			int(header.pop('''chunk_size'''       )       )
            SCREAMING_SNAKE_CASE__      :					List[str]							    =			list(header.values()       )

            SCREAMING_SNAKE_CASE__      :					List[Any]							    =			[]
            for i, presigned_url in enumerate(__a       ):
               with FileSlice(__a			,	seek_from=i * chunk_size			,	read_limit=__a       ) as data:
                  SCREAMING_SNAKE_CASE__      :					Optional[Any]							    =			requests.put(__a			,	data=__a       )
                  r.raise_for_status()
                  parts.append(
                      {
                          '''etag''': r.headers.get('''etag'''       ),
                          '''partNumber''': i + 1,
                      }       )
                  # In order to support progress reporting while data is uploading / downloading,
                  # the transfer process should post messages to stdout
                  write_msg(
                      {
                          '''event''': '''progress''',
                          '''oid''': oid,
                          '''bytesSoFar''': (i + 1) * chunk_size,
                          '''bytesSinceLast''': chunk_size,
                      }       )
                  # Not precise but that's ok.

            SCREAMING_SNAKE_CASE__      :					Optional[int]							    =			requests.post(
                __a			,	json={
                    '''oid''': oid,
                    '''parts''': parts,
                }			,	)
            r.raise_for_status()

            write_msg({'''event''': '''complete''', '''oid''': oid}       )









'''simple docstring'''






import unittest

from transformers import YosoConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
      import torch

      from transformers import (
          YosoForMaskedLM,
          YosoForMultipleChoice,
          YosoForQuestionAnswering,
          YosoForSequenceClassification,
          YosoForTokenClassification,
          YosoModel,
      )
      from transformers.models.yoso.modeling_yoso import YOSO_PRETRAINED_MODEL_ARCHIVE_LIST




class a_     :







   def __init__(self       , __a       , __a=1_3       , __a=7       , __a=True       , __a=True       , __a=True       , __a=True       , __a=9_9       , __a=3_2       , __a=5       , __a=4       , __a=3_7       , __a="gelu"       , __a=0.1       , __a=0.1       , __a=5_1_2       , __a=1_6       , __a=2       , __a=0.02       , __a=3       , __a=4       , __a=None       , )      ->   int:
      """simple docstring"""






      __snake_case					:   Union[str, Any]  				=			parent
      __snake_case					:   Dict  				=			batch_size
      __snake_case					:   Union[str, Any]  				=			seq_length
      __snake_case					:   str  				=			is_training
      __snake_case					:   Optional[Any]  				=			use_input_mask
      __snake_case					:   Dict  				=			use_token_type_ids
      __snake_case					:   Union[str, Any]  				=			use_labels
      __snake_case					:   Optional[Any]  				=			vocab_size
      __snake_case					:   Optional[Any]  				=			hidden_size
      __snake_case					:   int  				=			num_hidden_layers
      __snake_case					:   Any  				=			num_attention_heads
      __snake_case					:   Optional[Any]  				=			intermediate_size
      __snake_case					:   List[str]  				=			hidden_act
      __snake_case					:   List[Any]  				=			hidden_dropout_prob
      __snake_case					:   Optional[Any]  				=			attention_probs_dropout_prob
      __snake_case					:   str  				=			max_position_embeddings
      __snake_case					:   str  				=			type_vocab_size
      __snake_case					:   Dict  				=			type_sequence_label_size
      __snake_case					:   int  				=			initializer_range
      __snake_case					:   Dict  				=			num_labels
      __snake_case					:   Optional[int]  				=			num_choices
      __snake_case					:   Dict  				=			scope







   def      SCREAMING_SNAKE_CASE__       (self)      ->   List[Any]:
      """simple docstring"""






      __snake_case					:   Optional[int]  				=			ids_tensor([self.batch_size, self.seq_length]       , self.vocab_size)

      __snake_case					:   Optional[Any]  				=			None
      if self.use_input_mask:
         __snake_case					:   Union[str, Any]  				=			random_attention_mask([self.batch_size, self.seq_length])

      __snake_case					:   Tuple  				=			None
      if self.use_token_type_ids:
         __snake_case					:   Union[str, Any]  				=			ids_tensor([self.batch_size, self.seq_length]       , self.type_vocab_size)

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

      __snake_case					:   List[str]  				=			self.get_config()

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







   def      SCREAMING_SNAKE_CASE__       (self)      ->   List[str]:
      """simple docstring"""






      return YosoConfig(
          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       , is_decoder=__a       , initializer_range=self.initializer_range       , )







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Dict:
      """simple docstring"""






      __snake_case					:   Tuple  				=			self.get_config()
      __snake_case					:   Tuple  				=			3_0_0
      return config







   def      SCREAMING_SNAKE_CASE__       (self)      ->   str:
      """simple docstring"""






      (
          (
          __snake_case
      )				,(
          __snake_case
      )				,(
          __snake_case
      )				,(
          __snake_case
      )				,(
          __snake_case
      )				,(
          __snake_case
      )				,(
          __snake_case
      )				,
      )					:   Optional[Any]  				=			self.prepare_config_and_inputs()

      __snake_case					:   Optional[int]  				=			True
      __snake_case					:   Optional[int]  				=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
      __snake_case					:   Dict  				=			ids_tensor([self.batch_size, self.seq_length]       , vocab_size=2)

      return (
          config,
          input_ids,
          token_type_ids,
          input_mask,
          sequence_labels,
          token_labels,
          choice_labels,
          encoder_hidden_states,
          encoder_attention_mask,
      )







   def      SCREAMING_SNAKE_CASE__       (self       , __a       , __a       , __a       , __a       , __a       , __a       , __a)      ->   Dict:
      """simple docstring"""






      __snake_case					:   Optional[Any]  				=			YosoModel(config=__a)
      model.to(__a)
      model.eval()
      __snake_case					:   Optional[Any]  				=			model(__a       , attention_mask=__a       , token_type_ids=__a)
      __snake_case					:   int  				=			model(__a       , token_type_ids=__a)
      __snake_case					:   Optional[Any]  				=			model(__a)
      self.parent.assertEqual(result.last_hidden_state.shape       , (self.batch_size, self.seq_length, self.hidden_size))







   def      SCREAMING_SNAKE_CASE__       (self       , __a       , __a       , __a       , __a       , __a       , __a       , __a       , __a       , __a       , )      ->   int:
      """simple docstring"""






      __snake_case					:   Optional[Any]  				=			True
      __snake_case					:   Dict  				=			YosoModel(__a)
      model.to(__a)
      model.eval()
      __snake_case					:   List[str]  				=			model(
          __a       , attention_mask=__a       , token_type_ids=__a       , encoder_hidden_states=__a       , encoder_attention_mask=__a       , )
      __snake_case					:   List[Any]  				=			model(
          __a       , attention_mask=__a       , token_type_ids=__a       , encoder_hidden_states=__a       , )
      __snake_case					:   List[Any]  				=			model(__a       , attention_mask=__a       , token_type_ids=__a)
      self.parent.assertEqual(result.last_hidden_state.shape       , (self.batch_size, self.seq_length, self.hidden_size))







   def      SCREAMING_SNAKE_CASE__       (self       , __a       , __a       , __a       , __a       , __a       , __a       , __a)      ->   Union[str, Any]:
      """simple docstring"""






      __snake_case					:   Union[str, Any]  				=			YosoForMaskedLM(config=__a)
      model.to(__a)
      model.eval()
      __snake_case					:   Union[str, Any]  				=			model(__a       , attention_mask=__a       , token_type_ids=__a       , labels=__a)
      self.parent.assertEqual(result.logits.shape       , (self.batch_size, self.seq_length, self.vocab_size))







   def      SCREAMING_SNAKE_CASE__       (self       , __a       , __a       , __a       , __a       , __a       , __a       , __a)      ->   int:
      """simple docstring"""






      __snake_case					:   List[Any]  				=			YosoForQuestionAnswering(config=__a)
      model.to(__a)
      model.eval()
      __snake_case					:   str  				=			model(
          __a       , attention_mask=__a       , token_type_ids=__a       , start_positions=__a       , end_positions=__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      SCREAMING_SNAKE_CASE__       (self       , __a       , __a       , __a       , __a       , __a       , __a       , __a)      ->   int:
      """simple docstring"""






      __snake_case					:   str  				=			self.num_labels
      __snake_case					:   Dict  				=			YosoForSequenceClassification(__a)
      model.to(__a)
      model.eval()
      __snake_case					:   Optional[int]  				=			model(__a       , attention_mask=__a       , token_type_ids=__a       , labels=__a)
      self.parent.assertEqual(result.logits.shape       , (self.batch_size, self.num_labels))







   def      SCREAMING_SNAKE_CASE__       (self       , __a       , __a       , __a       , __a       , __a       , __a       , __a)      ->   Dict:
      """simple docstring"""






      __snake_case					:   str  				=			self.num_labels
      __snake_case					:   Dict  				=			YosoForTokenClassification(config=__a)
      model.to(__a)
      model.eval()
      __snake_case					:   Tuple  				=			model(__a       , attention_mask=__a       , token_type_ids=__a       , labels=__a)
      self.parent.assertEqual(result.logits.shape       , (self.batch_size, self.seq_length, self.num_labels))







   def      SCREAMING_SNAKE_CASE__       (self       , __a       , __a       , __a       , __a       , __a       , __a       , __a)      ->   str:
      """simple docstring"""






      __snake_case					:   int  				=			self.num_choices
      __snake_case					:   Tuple  				=			YosoForMultipleChoice(config=__a)
      model.to(__a)
      model.eval()
      __snake_case					:   Any  				=			input_ids.unsqueeze(1).expand(-1       , self.num_choices       , -1).contiguous()
      __snake_case					:   Tuple  				=			token_type_ids.unsqueeze(1).expand(-1       , self.num_choices       , -1).contiguous()
      __snake_case					:   List[str]  				=			input_mask.unsqueeze(1).expand(-1       , self.num_choices       , -1).contiguous()
      __snake_case					:   Optional[int]  				=			model(
          __a       , attention_mask=__a       , token_type_ids=__a       , labels=__a       , )
      self.parent.assertEqual(result.logits.shape       , (self.batch_size, self.num_choices))







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Dict:
      """simple docstring"""






      __snake_case					:   int  				=			self.prepare_config_and_inputs()
      (
          (
          __snake_case
      )				,(
          __snake_case
      )				,(
          __snake_case
      )				,(
          __snake_case
      )				,(
          __snake_case
      )				,(
          __snake_case
      )				,(
          __snake_case
      )				,
      )					:   Union[str, Any]  				=			config_and_inputs
      __snake_case					:   Optional[int]  				=			{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
      return config, inputs_dict




@require_torch
class a_     (     UpperCamelCase_		,   UpperCamelCase_		,   unittest.TestCase				):
   _snake_case	   =				(
       (
           YosoModel,
           YosoForMaskedLM,
           YosoForMultipleChoice,
           YosoForQuestionAnswering,
           YosoForSequenceClassification,
           YosoForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   _snake_case	   =				False
   _snake_case	   =				False
   _snake_case	   =				False

   _snake_case	   =				()
   _snake_case	   =				(
       {
           """feature-extraction""": YosoModel,
           """fill-mask""": YosoForMaskedLM,
           """question-answering""": YosoForQuestionAnswering,
           """text-classification""": YosoForSequenceClassification,
           """token-classification""": YosoForTokenClassification,
           """zero-shot""": YosoForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )







   def      SCREAMING_SNAKE_CASE__       (self)      ->   List[str]:
      """simple docstring"""






      __snake_case					:   Optional[Any]  				=			YosoModelTester(self)
      __snake_case					:   List[Any]  				=			ConfigTester(self       , config_class=__a       , hidden_size=3_7)







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Dict:
      """simple docstring"""






      self.config_tester.run_common_tests()







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Tuple:
      """simple docstring"""






      __snake_case					:   List[Any]  				=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__a)







   def      SCREAMING_SNAKE_CASE__       (self)      ->   List[str]:
      """simple docstring"""






      __snake_case					:   Tuple  				=			self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
         __snake_case					:   List[Any]  				=			type
         self.model_tester.create_and_check_model(*__a)







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Optional[Any]:
      """simple docstring"""






      __snake_case					:   List[Any]  				=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*__a)







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Optional[Any]:
      """simple docstring"""






      __snake_case					:   Optional[int]  				=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_multiple_choice(*__a)







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Union[str, Any]:
      """simple docstring"""






      __snake_case					:   str  				=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*__a)







   def      SCREAMING_SNAKE_CASE__       (self)      ->   int:
      """simple docstring"""






      __snake_case					:   Optional[int]  				=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*__a)







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Optional[int]:
      """simple docstring"""






      __snake_case					:   List[str]  				=			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*__a)







   @slow
   def      SCREAMING_SNAKE_CASE__       (self)      ->   List[str]:
      """simple docstring"""






      for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         __snake_case					:   Tuple  				=			YosoModel.from_pretrained(__a)
         self.assertIsNotNone(__a)







   def      SCREAMING_SNAKE_CASE__       (self)      ->   Any:
      """simple docstring"""






      return




@require_torch
class a_     (     unittest.TestCase				):







   @slow
   def      SCREAMING_SNAKE_CASE__       (self)      ->   Tuple:
      """simple docstring"""






      __snake_case					:   str  				=			YosoModel.from_pretrained('uw-madison/yoso-4096')
      __snake_case					:   str  				=			torch.tensor([[0, 1, 2, 3, 4, 5]])

      with torch.no_grad():
         __snake_case					:   str  				=			model(__a)[0]

      __snake_case					:   Union[str, Any]  				=			torch.Size((1, 6, 7_6_8))
      self.assertEqual(output.shape       , __a)

      __snake_case					:   Dict  				=			torch.tensor(
          [[[-0.0_611, 0.1_242, 0.0_840], [0.0_280, -0.0_048, 0.1_125], [0.0_106, 0.0_226, 0.0_751]]])

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







   @slow
   def      SCREAMING_SNAKE_CASE__       (self)      ->   int:
      """simple docstring"""






      __snake_case					:   Optional[Any]  				=			YosoForMaskedLM.from_pretrained('uw-madison/yoso-4096')
      __snake_case					:   Tuple  				=			torch.tensor([[0, 1, 2, 3, 4, 5]])

      with torch.no_grad():
         __snake_case					:   Any  				=			model(__a)[0]

      __snake_case					:   List[str]  				=			5_0_2_6_5

      __snake_case					:   List[str]  				=			torch.Size((1, 6, vocab_size))
      self.assertEqual(output.shape       , __a)

      __snake_case					:   str  				=			torch.tensor(
          [[[-2.1_313, -3.7_285, -2.2_407], [-2.7_047, -3.3_314, -2.6_408], [0.0_629, -2.5_166, -0.3_356]]])

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







   @slow
   def      SCREAMING_SNAKE_CASE__       (self)      ->   Optional[int]:
      """simple docstring"""






      __snake_case					:   Dict  				=			YosoForMaskedLM.from_pretrained('uw-madison/yoso-4096')
      __snake_case					:   Optional[Any]  				=			torch.arange(4_0_9_6).unsqueeze(0)

      with torch.no_grad():
         __snake_case					:   Dict  				=			model(__a)[0]

      __snake_case					:   List[str]  				=			5_0_2_6_5

      __snake_case					:   List[Any]  				=			torch.Size((1, 4_0_9_6, vocab_size))
      self.assertEqual(output.shape       , __a)

      __snake_case					:   Tuple  				=			torch.tensor(
          [[[-2.3_914, -4.3_742, -5.0_956], [-4.0_988, -4.2_384, -7.0_406], [-3.1_427, -3.7_192, -6.6_800]]])

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