
import collections
import inspect
import unittest
from typing import Dict, List, Tuple

from transformers import MaskFormerSwinConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, torch_device
from transformers.utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
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 MaskFormerSwinBackbone
			from transformers.models.maskformer import MaskFormerSwinModel




class     _SCREAMING_SNAKE_CASE     :



	def __init__(    self						,  lowercase						,  lowercase=13						,  lowercase=32						,  lowercase=2						,  lowercase=3						,  lowercase=16						,  lowercase=[1, 2, 1]						,  lowercase=[2, 2, 4]						,  lowercase=2						,  lowercase=2.0						,  lowercase=True						,  lowercase=0.0						,  lowercase=0.0						,  lowercase=0.1						,  lowercase="gelu"						,  lowercase=False						,  lowercase=True						,  lowercase=0.0_2						,  lowercase=1e-5						,  lowercase=True						,  lowercase=None						,  lowercase=True						,  lowercase=10						,  lowercase=8						,  lowercase=["stage1", "stage2", "stage3"]						,  lowercase=[1, 2, 3]						,  )				->       List[str]:
					lowerCamelCase_						=			parent
					lowerCamelCase_						=			batch_size
					lowerCamelCase_						=			image_size
					lowerCamelCase_						=			patch_size
					lowerCamelCase_						=			num_channels
					lowerCamelCase_						=			embed_dim
					lowerCamelCase_						=			depths
					lowerCamelCase_						=			num_heads
					lowerCamelCase_						=			window_size
					lowerCamelCase_						=			mlp_ratio
					lowerCamelCase_						=			qkv_bias
					lowerCamelCase_						=			hidden_dropout_prob
					lowerCamelCase_						=			attention_probs_dropout_prob
					lowerCamelCase_						=			drop_path_rate
					lowerCamelCase_						=			hidden_act
					lowerCamelCase_						=			use_absolute_embeddings
					lowerCamelCase_						=			patch_norm
					lowerCamelCase_						=			layer_norm_eps
					lowerCamelCase_						=			initializer_range
					lowerCamelCase_						=			is_training
					lowerCamelCase_						=			scope
					lowerCamelCase_						=			use_labels
					lowerCamelCase_						=			type_sequence_label_size
					lowerCamelCase_						=			encoder_stride
					lowerCamelCase_						=			out_features
					lowerCamelCase_						=			out_indices



	def     SCREAMING_SNAKE_CASE_(    self )				->       int:
					lowerCamelCase_						=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

					lowerCamelCase_						=			None
					if self.use_labels:
									lowerCamelCase_						=			ids_tensor([self.batch_size]						,  self.type_sequence_label_size )

					lowerCamelCase_						=			self.get_config()

					return config, pixel_values, labels



	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					return MaskFormerSwinConfig(
					    image_size=self.image_size						,  patch_size=self.patch_size						,  num_channels=self.num_channels						,  embed_dim=self.embed_dim						,  depths=self.depths						,  num_heads=self.num_heads						,  window_size=self.window_size						,  mlp_ratio=self.mlp_ratio						,  qkv_bias=self.qkv_bias						,  hidden_dropout_prob=self.hidden_dropout_prob						,  attention_probs_dropout_prob=self.attention_probs_dropout_prob						,  drop_path_rate=self.drop_path_rate						,  hidden_act=self.hidden_act						,  use_absolute_embeddings=self.use_absolute_embeddings						,  path_norm=self.patch_norm						,  layer_norm_eps=self.layer_norm_eps						,  initializer_range=self.initializer_range						,  encoder_stride=self.encoder_stride						,  out_features=self.out_features						,  out_indices=self.out_indices						,  )



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase						,  lowercase						,  lowercase )				->       List[Any]:
					lowerCamelCase_						=			MaskFormerSwinModel(config=lowercase )
					model.to(lowercase )
					model.eval()
					lowerCamelCase_						=			model(lowercase )

					lowerCamelCase_						=			((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths ) - 1))
					lowerCamelCase_						=			int(config.embed_dim * 2 ** (len(config.depths ) - 1) )

					self.parent.assertEqual(result.last_hidden_state.shape						,  (self.batch_size, expected_seq_len, expected_dim) )



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase						,  lowercase						,  lowercase )				->       Union[str, Any]:
					lowerCamelCase_						=			MaskFormerSwinBackbone(config=lowercase )
					model.to(lowercase )
					model.eval()
					lowerCamelCase_						=			model(lowercase )

					# verify feature maps
					self.parent.assertEqual(len(result.feature_maps )						,  len(config.out_features ) )
					self.parent.assertListEqual(list(result.feature_maps[0].shape )						,  [13, 16, 16, 16] )

					# verify channels
					self.parent.assertEqual(len(model.channels )						,  len(config.out_features ) )
					self.parent.assertListEqual(model.channels						,  [16, 32, 64] )

					# verify ValueError
					with self.parent.assertRaises(lowercase ):
									lowerCamelCase_						=			["stem"]
									lowerCamelCase_						=			MaskFormerSwinBackbone(config=lowercase )



	def     SCREAMING_SNAKE_CASE_(    self )				->       List[Any]:
					lowerCamelCase_						=			self.prepare_config_and_inputs()
					lowerCamelCase_	,    lowerCamelCase_	,    lowerCamelCase_						=			config_and_inputs
					lowerCamelCase_						=			{"pixel_values": pixel_values}
					return config, inputs_dict




@require_torch
class     _SCREAMING_SNAKE_CASE     (      snake_case_			,       snake_case_			,       unittest.TestCase						):
	lowerCAmelCase__         = (
	    (
	        MaskFormerSwinModel,
	        MaskFormerSwinBackbone,
	    )
	    if is_torch_available()
	    else ()
	)
	lowerCAmelCase__         = {'feature-extraction': MaskFormerSwinModel} if is_torch_available() else {}
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False
	lowerCAmelCase__         = False



	def     SCREAMING_SNAKE_CASE_(    self )				->       Tuple:
					lowerCamelCase_						=			MaskFormerSwinModelTester(self )
					lowerCamelCase_						=			ConfigTester(self						,  config_class=lowercase						,  embed_dim=37 )



	@require_torch_multi_gpu
	@unittest.skip(
	    reason=(
	        "`MaskFormerSwinModel` outputs `hidden_states_spatial_dimensions` which doesn't work well with"
	        " `nn.DataParallel`"
	    ) )
	def     SCREAMING_SNAKE_CASE_(    self )				->       List[Any]:
					pass



	def     SCREAMING_SNAKE_CASE_(    self )				->       List[str]:
					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     SCREAMING_SNAKE_CASE_(    self )				->       Tuple:
					return



	def     SCREAMING_SNAKE_CASE_(    self )				->       str:
					lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*lowercase )



	def     SCREAMING_SNAKE_CASE_(    self )				->       int:
					lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_backbone(*lowercase )



	@unittest.skip("Swin does not use inputs_embeds" )
	def     SCREAMING_SNAKE_CASE_(    self )				->       List[str]:
					pass



	@unittest.skip("Swin does not support feedforward chunking" )
	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					pass



	def     SCREAMING_SNAKE_CASE_(    self )				->       Optional[int]:
					lowerCamelCase_	,    lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs_for_common()

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



	def     SCREAMING_SNAKE_CASE_(    self )				->       Union[str, Any]:
					lowerCamelCase_	,    lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs_for_common()

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

									lowerCamelCase_						=			["pixel_values"]
									self.assertListEqual(arg_names[:1]						,  lowercase )



	@unittest.skip(reason="MaskFormerSwin is only used as backbone and doesn't support output_attentions" )
	def     SCREAMING_SNAKE_CASE_(    self )				->       int:
					pass



	@unittest.skip(reason="MaskFormerSwin is only used as an internal backbone" )
	def     SCREAMING_SNAKE_CASE_(    self )				->       List[Any]:
					pass



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase						,  lowercase						,  lowercase						,  lowercase )				->       List[Any]:
					lowerCamelCase_						=			model_class(lowercase )
					model.to(lowercase )
					model.eval()

					with torch.no_grad():
									lowerCamelCase_						=			model(**self._prepare_for_class(lowercase						,  lowercase ) )

					lowerCamelCase_						=			outputs.hidden_states

					lowerCamelCase_						=			getattr(
					    self.model_tester						,  "expected_num_hidden_layers"						,  len(self.model_tester.depths ) + 1 )
					self.assertEqual(len(lowercase )						,  lowercase )

					# Swin has a different seq_length
					lowerCamelCase_						=			(
					    config.patch_size
					    if isinstance(config.patch_size						,  collections.abc.Iterable )
					    else (config.patch_size, config.patch_size)
					)

					lowerCamelCase_						=			(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

					self.assertListEqual(
					    list(hidden_states[0].shape[-2:] )						,  [num_patches, self.model_tester.embed_dim]						,  )



	def     SCREAMING_SNAKE_CASE_(    self )				->       Dict:
					lowerCamelCase_	,    lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs_for_common()

					lowerCamelCase_						=			(
					    self.model_tester.image_size
					    if isinstance(self.model_tester.image_size						,  collections.abc.Iterable )
					    else (self.model_tester.image_size, self.model_tester.image_size)
					)

					for model_class in self.all_model_classes:
									lowerCamelCase_						=			True
									self.check_hidden_states_output(lowercase						,  lowercase						,  lowercase						,  lowercase )

									# check that output_hidden_states also work using config
									del inputs_dict["output_hidden_states"]
									lowerCamelCase_						=			True

									self.check_hidden_states_output(lowercase						,  lowercase						,  lowercase						,  lowercase )



	def     SCREAMING_SNAKE_CASE_(    self )				->       str:
					lowerCamelCase_	,    lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs_for_common()
					lowerCamelCase_						=			3

					lowerCamelCase_						=			(
					    self.model_tester.image_size
					    if isinstance(self.model_tester.image_size						,  collections.abc.Iterable )
					    else (self.model_tester.image_size, self.model_tester.image_size)
					)
					lowerCamelCase_						=			(
					    config.patch_size
					    if isinstance(config.patch_size						,  collections.abc.Iterable )
					    else (config.patch_size, config.patch_size)
					)

					lowerCamelCase_						=			image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
					lowerCamelCase_						=			image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

					for model_class in self.all_model_classes:
									lowerCamelCase_						=			True
									self.check_hidden_states_output(lowercase						,  lowercase						,  lowercase						,  (padded_height, padded_width) )

									# check that output_hidden_states also work using config
									del inputs_dict["output_hidden_states"]
									lowerCamelCase_						=			True
									self.check_hidden_states_output(lowercase						,  lowercase						,  lowercase						,  (padded_height, padded_width) )



	@unittest.skip(reason="MaskFormerSwin doesn't have pretrained checkpoints" )
	def     SCREAMING_SNAKE_CASE_(    self )				->       Dict:
					pass



	@unittest.skip(reason="This will be fixed once MaskFormerSwin is replaced by native Swin" )
	def     SCREAMING_SNAKE_CASE_(    self )				->       Dict:
					pass



	@unittest.skip(reason="This will be fixed once MaskFormerSwin is replaced by native Swin" )
	def     SCREAMING_SNAKE_CASE_(    self )				->       List[str]:
					pass



	def     SCREAMING_SNAKE_CASE_(    self )				->       List[Any]:
					lowerCamelCase_	,    lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs_for_common()

					def set_nan_tensor_to_zero(lowercase ):
									lowerCamelCase_						=			0
									return t

					def check_equivalence(lowercase						,  lowercase						,  lowercase						,  lowercase={} ):
									with torch.no_grad():
													lowerCamelCase_						=			model(**lowercase						,  return_dict=lowercase						,  **lowercase )
													lowerCamelCase_						=			model(**lowercase						,  return_dict=lowercase						,  **lowercase ).to_tuple()

													def recursive_check(lowercase						,  lowercase ):
																	if isinstance(lowercase						,  (List, Tuple) ):
																					for tuple_iterable_value, dict_iterable_value in zip(lowercase						,  lowercase ):
																									recursive_check(lowercase						,  lowercase )
																	elif isinstance(lowercase						,  lowercase ):
																					for tuple_iterable_value, dict_iterable_value in zip(
																					    tuple_object.values()						,  dict_object.values() ):
																									recursive_check(lowercase						,  lowercase )
																	elif tuple_object is None:
																					return
																	else:
																					self.assertTrue(
																					    torch.allclose(
																					        set_nan_tensor_to_zero(lowercase )						,  set_nan_tensor_to_zero(lowercase )						,  atol=1e-5 )						,  msg=(
																					        "Tuple and dict output are not equal. Difference:"
																					        f' {torch.max(torch.abs(tuple_object - dict_object ) )}. Tuple has `nan`:'
																					        f' {torch.isnan(lowercase ).any()} and `inf`: {torch.isinf(lowercase )}. Dict has'
																					        f' `nan`: {torch.isnan(lowercase ).any()} and `inf`: {torch.isinf(lowercase )}.'
																					    )						,  )

													recursive_check(lowercase						,  lowercase )

					for model_class in self.all_model_classes:
									lowerCamelCase_						=			model_class(lowercase )
									model.to(lowercase )
									model.eval()

									lowerCamelCase_						=			self._prepare_for_class(lowercase						,  lowercase )
									lowerCamelCase_						=			self._prepare_for_class(lowercase						,  lowercase )
									check_equivalence(lowercase						,  lowercase						,  lowercase )

									lowerCamelCase_						=			self._prepare_for_class(lowercase						,  lowercase						,  return_labels=lowercase )
									lowerCamelCase_						=			self._prepare_for_class(lowercase						,  lowercase						,  return_labels=lowercase )
									check_equivalence(lowercase						,  lowercase						,  lowercase )

									lowerCamelCase_						=			self._prepare_for_class(lowercase						,  lowercase )
									lowerCamelCase_						=			self._prepare_for_class(lowercase						,  lowercase )
									check_equivalence(lowercase						,  lowercase						,  lowercase						,  {"output_hidden_states": True} )

									lowerCamelCase_						=			self._prepare_for_class(lowercase						,  lowercase						,  return_labels=lowercase )
									lowerCamelCase_						=			self._prepare_for_class(lowercase						,  lowercase						,  return_labels=lowercase )
									check_equivalence(lowercase						,  lowercase						,  lowercase						,  {"output_hidden_states": True} )




@require_torch
class     _SCREAMING_SNAKE_CASE     (      unittest.TestCase			,       snake_case_						):
	lowerCAmelCase__         = (MaskFormerSwinBackbone,) if is_torch_available() else ()
	lowerCAmelCase__         = MaskFormerSwinConfig



	def     SCREAMING_SNAKE_CASE_(    self )				->       Any:
					lowerCamelCase_						=			MaskFormerSwinModelTester(self )



	def     SCREAMING_SNAKE_CASE_(    self )				->       List[str]:
					lowerCamelCase_	,    lowerCamelCase_						=			self.model_tester.prepare_config_and_inputs_for_common()
					lowerCamelCase_						=			inputs_dict["pixel_values"].shape[0]

					for backbone_class in self.all_model_classes:
									lowerCamelCase_						=			backbone_class(lowercase )
									backbone.to(lowercase )
									backbone.eval()

									lowerCamelCase_						=			backbone(**lowercase )

									# Test default outputs and verify feature maps
									self.assertIsInstance(outputs.feature_maps						,  lowercase )
									self.assertTrue(len(outputs.feature_maps ) == len(backbone.channels ) )
									for feature_map, n_channels in zip(outputs.feature_maps						,  backbone.channels ):
													self.assertTrue(feature_map.shape[:2]						,  (batch_size, n_channels) )
									self.assertIsNone(outputs.hidden_states )
									self.assertIsNone(outputs.attentions )

									# Test output_hidden_states=True
									lowerCamelCase_						=			backbone(**lowercase						,  output_hidden_states=lowercase )
									self.assertIsNotNone(outputs.hidden_states )
									self.assertTrue(len(outputs.hidden_states )						,  len(backbone.stage_names ) )
									# We skip the stem layer
									for hidden_states, n_channels in zip(outputs.hidden_states[1:]						,  backbone.channels ):
													for hidden_state in hidden_states:
																	# Hidden states are in the format (batch_size, (height * width), n_channels)
																	lowerCamelCase_	,    lowerCamelCase_	,    lowerCamelCase_						=			hidden_state.shape
																	self.assertTrue((h_batch_size, h_n_channels)						,  (batch_size, n_channels) )

            # Test output_attentions=True
									if self.has_attentions:
													lowerCamelCase_						=			backbone(**lowercase						,  output_attentions=lowercase )
													self.assertIsNotNone(outputs.attentions )










"""simple docstring"""


import unittest

import numpy as np

from transformers import RoFormerConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
       import jax.numpy as jnp

       from transformers.models.roformer.modeling_flax_roformer import (
           FlaxRoFormerForMaskedLM,
           FlaxRoFormerForMultipleChoice,
           FlaxRoFormerForQuestionAnswering,
           FlaxRoFormerForSequenceClassification,
           FlaxRoFormerForTokenClassification,
           FlaxRoFormerModel,
       )


class 					snake_case		(							unittest.TestCase						):


   '''simple docstring'''

   def __init__(	self :     Optional[Any],       _lowerCamelCase :     Tuple,       _lowerCamelCase :     List[str]=13,       _lowerCamelCase :     Optional[Any]=7,       _lowerCamelCase :     Optional[int]=True,       _lowerCamelCase :     int=True,       _lowerCamelCase :     List[str]=True,       _lowerCamelCase :     Optional[Any]=True,       _lowerCamelCase :     int=99,       _lowerCamelCase :     Optional[int]=32,       _lowerCamelCase :     Tuple=5,       _lowerCamelCase :     Tuple=4,       _lowerCamelCase :     str=37,       _lowerCamelCase :     Union[str, Any]="gelu",       _lowerCamelCase :     int=0.1,       _lowerCamelCase :     List[Any]=0.1,       _lowerCamelCase :     Dict=5_12,       _lowerCamelCase :     List[Any]=16,       _lowerCamelCase :     Any=2,       _lowerCamelCase :     Any=0.02,       _lowerCamelCase :     Dict=4,       ):

     '''simple docstring'''
     __A									= parent
     __A									= batch_size
     __A									= seq_length
     __A									= is_training
     __A									= use_attention_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_choices


   def 				_SCREAMING_SNAKE_CASE				(	self :     Any  ):

     '''simple docstring'''
     __A									= ids_tensor([self.batch_size, self.seq_length],       self.vocab_size  )

     __A									= None
     if self.use_attention_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									= RoFormerConfig(
         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=_lowerCamelCase,       initializer_range=self.initializer_range,       )

     return config, input_ids, token_type_ids, attention_mask






   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[int]  ):

     '''simple docstring'''
     __A									= self.prepare_config_and_inputs()
     __A			,    __A			,    __A			,    __A									= config_and_inputs
     __A									= {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
     return config, inputs_dict





@require_flax
class 					snake_case		(							_lowerCAmelCase   ,	unittest.TestCase						):


   '''simple docstring'''

   A_						:	Dict						 = True

   A_						:	Tuple						 = (
       (
           FlaxRoFormerModel,
           FlaxRoFormerForMaskedLM,
           FlaxRoFormerForSequenceClassification,
           FlaxRoFormerForTokenClassification,
           FlaxRoFormerForMultipleChoice,
           FlaxRoFormerForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )


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

     '''simple docstring'''
     __A									= FlaxRoFormerModelTester(self  )






   @slow
   def 				_SCREAMING_SNAKE_CASE				(	self :     Any  ):

     '''simple docstring'''
     for model_class_name in self.all_model_classes:
       __A									= model_class_name.from_pretrained('''junnyu/roformer_chinese_small''',       from_pt=_lowerCamelCase  )
       __A									= model(np.ones((1, 1)  )  )
       self.assertIsNotNone(_lowerCamelCase  )





@require_flax
class 					snake_case		(							unittest.TestCase						):


   '''simple docstring'''

   @slow
   def 				_SCREAMING_SNAKE_CASE				(	self :     Dict  ):

     '''simple docstring'''
     __A									= FlaxRoFormerForMaskedLM.from_pretrained('''junnyu/roformer_chinese_base'''  )
     __A									= jnp.array([[0, 1, 2, 3, 4, 5]]  )
     __A									= model(_lowerCamelCase  )[0]

     __A									= 5_00_00

     __A									= (1, 6, vocab_size)
     self.assertEqual(output.shape,       _lowerCamelCase  )

     __A									= jnp.array(
         [[[-0.12_05, -1.02_65, 0.29_22], [-1.51_34, 0.19_74, 0.15_19], [-5.01_35, -3.90_03, -0.84_04]]]  )

     self.assertTrue(jnp.allclose(output[:, :3, :3],       _lowerCamelCase,       atol=1e-4  )  )




