






'''simple docstring'''







from __future__ import annotations

import unittest

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

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


if is_tf_available():
					import numpy
					import tensorflow as tf

					from transformers.models.esm.modeling_tf_esm import (
					    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
					    TFEsmForMaskedLM,
					    TFEsmForSequenceClassification,
					    TFEsmForTokenClassification,
					    TFEsmModel,
					)


class 		__UpperCAmelCase   :


							def __init__( self					,     lowerCAmelCase_					,     ):





								"""simple docstring"""






								_snake_case     		=	parent
								_snake_case     		=	13
								_snake_case     		=	7
								_snake_case     		=	True
								_snake_case     		=	True
								_snake_case     		=	True
								_snake_case     		=	99
								_snake_case     		=	32
								_snake_case     		=	2
								_snake_case     		=	4
								_snake_case     		=	37
								_snake_case     		=	'gelu'
								_snake_case     		=	0.1
								_snake_case     		=	0.1
								_snake_case     		=	5_12
								_snake_case     		=	16
								_snake_case     		=	2
								_snake_case     		=	0.02
								_snake_case     		=	3
								_snake_case     		=	4
								_snake_case     		=	None


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	ids_tensor([self.batch_size, self.seq_length]					,     self.vocab_size						)

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

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

								_snake_case     		=	EsmConfig(
								    vocab_size=self.vocab_size					,     hidden_size=self.hidden_size					,     num_hidden_layers=self.num_hidden_layers					,     pad_token_id=1					,     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					,     initializer_range=self.initializer_range					,     )

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


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								(
								    (
								    _snake_case      
								)		,				(
								    _snake_case      
								)		,				(
								    _snake_case      
								)		,				(
								    _snake_case      
								)		,				(
								    _snake_case      
								)		,				(
								    _snake_case      
								)		,				      
								)     		=	self.prepare_config_and_inputs()

								_snake_case     		=	True
								_snake_case     		=	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
								_snake_case     		=	ids_tensor([self.batch_size, self.seq_length]					,     vocab_size=2						)

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


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	TFEsmModel(config=A_						)
								_snake_case     		=	{'input_ids': input_ids, 'attention_mask': input_mask}
								_snake_case     		=	model(A_						)

								_snake_case     		=	[input_ids, input_mask]
								_snake_case     		=	model(A_						)

								_snake_case     		=	model(A_						)

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


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     ):





								"""simple docstring"""






								_snake_case     		=	True

								_snake_case     		=	TFEsmModel(config=A_						)
								_snake_case     		=	{
								    'input_ids': input_ids,
								    'attention_mask': input_mask,
								    'encoder_hidden_states': encoder_hidden_states,
								    'encoder_attention_mask': encoder_attention_mask,
								}
								_snake_case     		=	model(A_						)

								_snake_case     		=	[input_ids, input_mask]
								_snake_case     		=	model(A_					,     encoder_hidden_states=A_						)

								# Also check the case where encoder outputs are not passed
								_snake_case     		=	model(A_					,     attention_mask=A_						)

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


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	TFEsmForMaskedLM(config=A_						)
								_snake_case     		=	model([input_ids, input_mask]						)
								self.parent.assertEqual(result.logits.shape					,     (self.batch_size, self.seq_length, self.vocab_size)						)


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	self.num_labels
								_snake_case     		=	TFEsmForTokenClassification(config=A_						)
								_snake_case     		=	{'input_ids': input_ids, 'attention_mask': input_mask}
								_snake_case     		=	model(A_						)
								self.parent.assertEqual(result.logits.shape					,     (self.batch_size, self.seq_length, self.num_labels)						)




							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.prepare_config_and_inputs()
								(
								    (
								    _snake_case      
								)		,				(
								    _snake_case      
								)		,				(
								    _snake_case      
								)		,				(
								    _snake_case      
								)		,				(
								    _snake_case      
								)		,				(
								    _snake_case      
								)		,				      
								)     		=	config_and_inputs
								_snake_case     		=	{'input_ids': input_ids, 'attention_mask': input_mask}
								return config, inputs_dict


@require_tf
class 		__UpperCAmelCase   (      _lowercase		,     _lowercase		,     unittest.TestCase		):
							__lowercase   	= (
							    (
							        TFEsmModel,
							        TFEsmForMaskedLM,
							        TFEsmForSequenceClassification,
							        TFEsmForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							__lowercase   	= (
							    {
							        '''feature-extraction''': TFEsmModel,
							        '''fill-mask''': TFEsmForMaskedLM,
							        '''text-classification''': TFEsmForSequenceClassification,
							        '''token-classification''': TFEsmForTokenClassification,
							        '''zero-shot''': TFEsmForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							__lowercase   	= False
							__lowercase   	= False


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	TFEsmModelTester(self						)
								_snake_case     		=	ConfigTester(self					,     config_class=A_					,     hidden_size=37						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								self.config_tester.run_common_tests()


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*A_						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(*A_						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*A_						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*A_						)


							@slow
							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									_snake_case     		=	TFEsmModel.from_pretrained(A_						)
									self.assertIsNotNone(A_						)


							@unittest.skip('Protein models do not support embedding resizing.'						)
							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								pass


							@unittest.skip('Protein models do not support embedding resizing.'						)
							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								pass




							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case		,				_snake_case     		=	self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
									_snake_case     		=	model_class(A_						)
									assert isinstance(model.get_input_embeddings()					,     tf.keras.layers.Layer						)
									if model_class is TFEsmForMaskedLM:
										# Output embedding test differs from the main test because they're a matrix, not a layer
										_snake_case     		=	model.get_bias()
										assert isinstance(A_					,     A_						)
										for k, v in name.items():
											assert isinstance(A_					,     tf.Variable						)
									else:
										_snake_case     		=	model.get_output_embeddings()
										assert x is None
										_snake_case     		=	model.get_bias()
										assert name is None


@require_tf
class 		__UpperCAmelCase   (      unittest.TestCase		):


							@slow
							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'						)

								_snake_case     		=	tf.constant([[0, 1, 2, 3, 4, 5]]						)
								_snake_case     		=	model(A_						)[0]
								_snake_case     		=	[1, 6, 33]
								self.assertEqual(list(output.numpy().shape						)					,     A_						)
								# compare the actual values for a slice.
								_snake_case     		=	tf.constant(
								    [
								        [
								            [8.921518, -10.589814, -6.4671307],
								            [-6.3967156, -13.911377, -1.1211915],
								            [-7.781247, -13.951557, -3.740592],
								        ]
								    ]						)
								self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()					,     expected_slice.numpy()					,     atol=1E-2						)						)




							@slow
							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'						)

								_snake_case     		=	tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
								_snake_case     		=	model(A_						)[0]
								# compare the actual values for a slice.
								_snake_case     		=	tf.constant(
								    [
								        [
								            [0.14443092, 0.54125327, 0.3247739],
								            [0.30340484, 0.00526676, 0.31077722],
								            [0.32278043, -0.24987096, 0.3414628],
								        ]
								    ]						)
								self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()					,     expected_slice.numpy()					,     atol=1E-4						)						)



'''simple docstring'''







from __future__ import annotations

import math







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			list:
	if len(__A			) != 2 or len(a[0]			) != 2 or len(__A			) != 2 or len(b[0]			) != 2:
		raise Exception('Matrices are not 2x2'			)
	_snake_case     		=	[
	    [a[0][0] * b[0][0] + a[0][1] * b[1][0], a[0][0] * b[0][1] + a[0][1] * b[1][1]],
	    [a[1][0] * b[0][0] + a[1][1] * b[1][0], a[1][0] * b[0][1] + a[1][1] * b[1][1]],
	]
	return new_matrix







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			Dict:
	return [
	    [matrix_a[row][col] + matrix_b[row][col] for col in range(len(matrix_a[row]			)			)]
	    for row in range(len(__A			)			)
	]







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			Dict:
	return [
	    [matrix_a[row][col] - matrix_b[row][col] for col in range(len(matrix_a[row]			)			)]
	    for row in range(len(__A			)			)
	]







def        SCREAMING_SNAKE_CASE__			(     __A			) ->			tuple[list, list, list, list]:
	if len(__A			) % 2 != 0 or len(a[0]			) % 2 != 0:
		raise Exception('Odd matrices are not supported!'			)

	_snake_case     		=	len(__A			)
	_snake_case     		=	matrix_length // 2

	_snake_case     		=	[[a[i][j] for j in range(__A    ,       __A			)] for i in range(__A			)]
	_snake_case     		=	[
	    [a[i][j] for j in range(__A    ,       __A			)] for i in range(__A    ,       __A			)
	]

	_snake_case     		=	[[a[i][j] for j in range(__A			)] for i in range(__A			)]
	_snake_case     		=	[[a[i][j] for j in range(__A			)] for i in range(__A    ,       __A			)]

	return top_left, top_right, bot_left, bot_right







def        SCREAMING_SNAKE_CASE__			(     __A			) ->			tuple[int, int]:
	return len(__A			), len(matrix[0]			)







def        SCREAMING_SNAKE_CASE__			(     __A			) ->			None:
	print('\n'.join(str(__A			) for line in matrix			)			)







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			list:
	if matrix_dimensions(__A			) == (2, 2):
		return default_matrix_multiplication(__A    ,       __A			)

	_snake_case		,				_snake_case		,				_snake_case		,				_snake_case     		=	split_matrix(__A			)
	_snake_case		,				_snake_case		,				_snake_case		,				_snake_case     		=	split_matrix(__A			)

	_snake_case     		=	actual_strassen(__A    ,       matrix_subtraction(__A    ,       __A			)			)
	_snake_case     		=	actual_strassen(matrix_addition(__A    ,       __A			)    ,       __A			)
	_snake_case     		=	actual_strassen(matrix_addition(__A    ,       __A			)    ,       __A			)
	_snake_case     		=	actual_strassen(__A    ,       matrix_subtraction(__A    ,       __A			)			)
	_snake_case     		=	actual_strassen(matrix_addition(__A    ,       __A			)    ,       matrix_addition(__A    ,       __A			)			)
	_snake_case     		=	actual_strassen(matrix_subtraction(__A    ,       __A			)    ,       matrix_addition(__A    ,       __A			)			)
	_snake_case     		=	actual_strassen(matrix_subtraction(__A    ,       __A			)    ,       matrix_addition(__A    ,       __A			)			)

	_snake_case     		=	matrix_addition(matrix_subtraction(matrix_addition(__A    ,       __A			)    ,       __A			)    ,       __A			)
	_snake_case     		=	matrix_addition(__A    ,       __A			)
	_snake_case     		=	matrix_addition(__A    ,       __A			)
	_snake_case     		=	matrix_subtraction(matrix_subtraction(matrix_addition(__A    ,       __A			)    ,       __A			)    ,       __A			)

	# construct the new matrix from our 4 quadrants
	_snake_case     		=	[]
	for i in range(len(__A			)			):
		new_matrix.append(top_left[i] + top_right[i]			)
	for i in range(len(__A			)			):
		new_matrix.append(bot_left[i] + bot_right[i]			)
	return new_matrix







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			list:
	if matrix_dimensions(__A			)[1] != matrix_dimensions(__A			)[0]:
		_snake_case     		=	(
		    'Unable to multiply these matrices, please check the dimensions.\n'
		    F'Matrix A: {matrixa}\n'
		    F'Matrix B: {matrixa}'
		)
		raise Exception(__A			)
	_snake_case     		=	matrix_dimensions(__A			)
	_snake_case     		=	matrix_dimensions(__A			)

	if dimensiona[0] == dimensiona[1] and dimensiona[0] == dimensiona[1]:
		return [matrixa, matrixa]

	_snake_case     		=	max(*__A    ,       *__A			)
	_snake_case     		=	int(math.pow(2    ,       math.ceil(math.loga(__A			)			)			)			)
	_snake_case     		=	matrixa
	_snake_case     		=	matrixa

	# Adding zeros to the matrices so that the arrays dimensions are the same and also
	# power of 2
	for i in range(0    ,       __A			):
		if i < dimensiona[0]:
			for _ in range(dimensiona[1]    ,       __A			):
				new_matrixa[i].append(0			)
		else:
			new_matrixa.append([0] * maxim			)
		if i < dimensiona[0]:
			for _ in range(dimensiona[1]    ,       __A			):
				new_matrixa[i].append(0			)
		else:
			new_matrixa.append([0] * maxim			)

	_snake_case     		=	actual_strassen(__A    ,       __A			)

	# Removing the additional zeros
	for i in range(0    ,       __A			):
		if i < dimensiona[0]:
			for _ in range(dimensiona[1]    ,       __A			):
				final_matrix[i].pop()
		else:
			final_matrix.pop()
	return final_matrix


if __name__ == "__main__":
					lowercase					:							Any   =      [
					    [2, 3, 4, 5],
					    [6, 4, 3, 1],
					    [2, 3, 6, 7],
					    [3, 1, 2, 4],
					    [2, 3, 4, 5],
					    [6, 4, 3, 1],
					    [2, 3, 6, 7],
					    [3, 1, 2, 4],
					    [2, 3, 4, 5],
					    [6, 2, 3, 1],
					]
					lowercase					:							List[str]   =      [[0, 2, 1, 1], [16, 2, 3, 3], [2, 2, 7, 7], [13, 11, 22, 4]]
					print(strassen(matrixa, matrixa))

