

'''simple docstring'''
from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)
@flax.struct.dataclass
class 					_snake_case		(__SCREAMING_SNAKE_CASE):
			__A  :	jnp.ndarray
			__A  :	jnp.ndarray




class 					_snake_case		(nn.Module):
			__A  :	int
			__A  :	Tuple[int]        =(16, 32, 96, 2_56)
			__A  :	jnp.dtype        =jnp.floataa






			def 		UpperCamelCase__							(  self     ):
				UpperCAmelCase_     :      Tuple						=  nn.Conv(
				    self.block_out_channels[0]			,kernel_size=(3, 3)			,padding=((1, 1), (1, 1))			,dtype=self.dtype			,)

				UpperCAmelCase_     :      Any						=  []
				for i in range(len(self.block_out_channels     ) - 1     ):
					UpperCAmelCase_     :      str						=  self.block_out_channels[i]
					UpperCAmelCase_     :      Tuple						=  self.block_out_channels[i + 1]
					UpperCAmelCase_     :      Any						=  nn.Conv(
					    _snake_case			,kernel_size=(3, 3)			,padding=((1, 1), (1, 1))			,dtype=self.dtype			,)
					blocks.append(_snake_case     )
					UpperCAmelCase_     :      Optional[int]						=  nn.Conv(
					    _snake_case			,kernel_size=(3, 3)			,strides=(2, 2)			,padding=((1, 1), (1, 1))			,dtype=self.dtype			,)
					blocks.append(_snake_case     )
				UpperCAmelCase_     :      Tuple						=  blocks

				UpperCAmelCase_     :      str						=  nn.Conv(
				    self.conditioning_embedding_channels			,kernel_size=(3, 3)			,padding=((1, 1), (1, 1))			,kernel_init=nn.initializers.zeros_init()			,bias_init=nn.initializers.zeros_init()			,dtype=self.dtype			,)






			def __call__(  self			,_snake_case     ):
				UpperCAmelCase_     :      Dict						=  self.conv_in(_snake_case     )
				UpperCAmelCase_     :      int						=  nn.silu(_snake_case     )

				for block in self.blocks:
					UpperCAmelCase_     :      Union[str, Any]						=  block(_snake_case     )
					UpperCAmelCase_     :      Tuple						=  nn.silu(_snake_case     )

				UpperCAmelCase_     :      Union[str, Any]						=  self.conv_out(_snake_case     )

				return embedding




@flax_register_to_config
class 					_snake_case		(nn.Module			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE):
			__A  :	int        =32
			__A  :	int        =4
			__A  :	Tuple[str]        =(
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "CrossAttnDownBlock2D",
			    "DownBlock2D",
			)
			__A  :	Union[bool, Tuple[bool]]        =False
			__A  :	Tuple[int]        =(3_20, 6_40, 12_80, 12_80)
			__A  :	int        =2
			__A  :	Union[int, Tuple[int]]        =8
			__A  :	Optional[Union[int, Tuple[int]]]        =None
			__A  :	int        =12_80
			__A  :	float        =0.0
			__A  :	bool        =False
			__A  :	jnp.dtype        =jnp.floataa
			__A  :	bool        =True
			__A  :	int        =0
			__A  :	str        ="rgb"
			__A  :	Tuple[int]        =(16, 32, 96, 2_56)






			def 		UpperCamelCase__							(  self			,_snake_case     ):
				# init input tensors
				UpperCAmelCase_     :      List[Any]						=  (1, self.in_channels, self.sample_size, self.sample_size)
				UpperCAmelCase_     :      Any						=  jnp.zeros(_snake_case			,dtype=jnp.floataa     )
				UpperCAmelCase_     :      Tuple						=  jnp.ones((1,)			,dtype=jnp.intaa     )
				UpperCAmelCase_     :      List[Any]						=  jnp.zeros((1, 1, self.cross_attention_dim)			,dtype=jnp.floataa     )
				UpperCAmelCase_     :      Union[str, Any]						=  (1, 3, self.sample_size * 8, self.sample_size * 8)
				UpperCAmelCase_     :      int						=  jnp.zeros(_snake_case			,dtype=jnp.floataa     )

				UpperCAmelCase_					,					UpperCAmelCase_     :      str						=  jax.random.split(_snake_case     )
				UpperCAmelCase_     :      Optional[int]						=  {"params": params_rng, "dropout": dropout_rng}

				return self.init(_snake_case			,_snake_case			,_snake_case			,_snake_case			,_snake_case     )["params"]






			def 		UpperCamelCase__							(  self     ):
				UpperCAmelCase_     :      Tuple						=  self.block_out_channels
				UpperCAmelCase_     :      Union[str, Any]						=  block_out_channels[0] * 4

				# If `num_attention_heads` is not defined (which is the case for most models)
				# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
				# The reason for this behavior is to correct for incorrectly named variables that were introduced
				# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
				# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
				# which is why we correct for the naming here.
				UpperCAmelCase_     :      Dict						=  self.num_attention_heads or self.attention_head_dim

				# input
				UpperCAmelCase_     :      str						=  nn.Conv(
				    block_out_channels[0]			,kernel_size=(3, 3)			,strides=(1, 1)			,padding=((1, 1), (1, 1))			,dtype=self.dtype			,)

				# time
				UpperCAmelCase_     :      Dict						=  FlaxTimesteps(
				    block_out_channels[0]			,flip_sin_to_cos=self.flip_sin_to_cos			,freq_shift=self.config.freq_shift     )
				UpperCAmelCase_     :      List[Any]						=  FlaxTimestepEmbedding(_snake_case			,dtype=self.dtype     )

				UpperCAmelCase_     :      int						=  FlaxControlNetConditioningEmbedding(
				    conditioning_embedding_channels=block_out_channels[0]			,block_out_channels=self.conditioning_embedding_out_channels			,)

				UpperCAmelCase_     :      Any						=  self.only_cross_attention
				if isinstance(_snake_case			,_snake_case     ):
					UpperCAmelCase_     :      Union[str, Any]						=  (only_cross_attention,) * len(self.down_block_types     )

				if isinstance(_snake_case			,_snake_case     ):
					UpperCAmelCase_     :      Tuple						=  (num_attention_heads,) * len(self.down_block_types     )

				# down
				UpperCAmelCase_     :      List[str]						=  []
				UpperCAmelCase_     :      str						=  []

				UpperCAmelCase_     :      List[str]						=  block_out_channels[0]

				UpperCAmelCase_     :      Union[str, Any]						=  nn.Conv(
				    _snake_case			,kernel_size=(1, 1)			,padding="VALID"			,kernel_init=nn.initializers.zeros_init()			,bias_init=nn.initializers.zeros_init()			,dtype=self.dtype			,)
				controlnet_down_blocks.append(_snake_case     )

				for i, down_block_type in enumerate(self.down_block_types     ):
					UpperCAmelCase_     :      Tuple						=  output_channel
					UpperCAmelCase_     :      List[Any]						=  block_out_channels[i]
					UpperCAmelCase_     :      int						=  i == len(_snake_case     ) - 1

					if down_block_type == "CrossAttnDownBlock2D":
						UpperCAmelCase_     :      int						=  FlaxCrossAttnDownBlockaD(
						    in_channels=_snake_case			,out_channels=_snake_case			,dropout=self.dropout			,num_layers=self.layers_per_block			,num_attention_heads=num_attention_heads[i]			,add_downsample=not is_final_block			,use_linear_projection=self.use_linear_projection			,only_cross_attention=only_cross_attention[i]			,dtype=self.dtype			,)
					else:
						UpperCAmelCase_     :      List[Any]						=  FlaxDownBlockaD(
						    in_channels=_snake_case			,out_channels=_snake_case			,dropout=self.dropout			,num_layers=self.layers_per_block			,add_downsample=not is_final_block			,dtype=self.dtype			,)

					down_blocks.append(_snake_case     )

					for _ in range(self.layers_per_block     ):
						UpperCAmelCase_     :      List[Any]						=  nn.Conv(
						    _snake_case			,kernel_size=(1, 1)			,padding="VALID"			,kernel_init=nn.initializers.zeros_init()			,bias_init=nn.initializers.zeros_init()			,dtype=self.dtype			,)
						controlnet_down_blocks.append(_snake_case     )

					if not is_final_block:
						UpperCAmelCase_     :      int						=  nn.Conv(
						    _snake_case			,kernel_size=(1, 1)			,padding="VALID"			,kernel_init=nn.initializers.zeros_init()			,bias_init=nn.initializers.zeros_init()			,dtype=self.dtype			,)
						controlnet_down_blocks.append(_snake_case     )

				UpperCAmelCase_     :      int						=  down_blocks
				UpperCAmelCase_     :      Union[str, Any]						=  controlnet_down_blocks

				# mid
				UpperCAmelCase_     :      int						=  block_out_channels[-1]
				UpperCAmelCase_     :      Tuple						=  FlaxUNetMidBlockaDCrossAttn(
				    in_channels=_snake_case			,dropout=self.dropout			,num_attention_heads=num_attention_heads[-1]			,use_linear_projection=self.use_linear_projection			,dtype=self.dtype			,)

				UpperCAmelCase_     :      List[str]						=  nn.Conv(
				    _snake_case			,kernel_size=(1, 1)			,padding="VALID"			,kernel_init=nn.initializers.zeros_init()			,bias_init=nn.initializers.zeros_init()			,dtype=self.dtype			,)






			def __call__(  self			,_snake_case			,_snake_case			,_snake_case			,_snake_case			,_snake_case = 1.0			,_snake_case = True			,_snake_case = False			,):
				UpperCAmelCase_     :      int						=  self.controlnet_conditioning_channel_order
				if channel_order == "bgr":
					UpperCAmelCase_     :      Union[str, Any]						=  jnp.flip(_snake_case			,axis=1     )

				# 1. time
				if not isinstance(_snake_case			,jnp.ndarray     ):
					UpperCAmelCase_     :      Optional[int]						=  jnp.array([timesteps]			,dtype=jnp.intaa     )
				elif isinstance(_snake_case			,jnp.ndarray     ) and len(timesteps.shape     ) == 0:
					UpperCAmelCase_     :      str						=  timesteps.astype(dtype=jnp.floataa     )
					UpperCAmelCase_     :      Optional[int]						=  jnp.expand_dims(_snake_case			,0     )

				UpperCAmelCase_     :      str						=  self.time_proj(_snake_case     )
				UpperCAmelCase_     :      Optional[Any]						=  self.time_embedding(_snake_case     )

				# 2. pre-process
				UpperCAmelCase_     :      Union[str, Any]						=  jnp.transpose(_snake_case			,(0, 2, 3, 1)     )
				UpperCAmelCase_     :      List[str]						=  self.conv_in(_snake_case     )

				UpperCAmelCase_     :      Tuple						=  jnp.transpose(_snake_case			,(0, 2, 3, 1)     )
				UpperCAmelCase_     :      Optional[int]						=  self.controlnet_cond_embedding(_snake_case     )
				sample += controlnet_cond

				# 3. down
				UpperCAmelCase_     :      Union[str, Any]						=  (sample,)
				for down_block in self.down_blocks:
					if isinstance(_snake_case			,_snake_case     ):
						UpperCAmelCase_					,					UpperCAmelCase_     :      Union[str, Any]						=  down_block(_snake_case			,_snake_case			,_snake_case			,deterministic=not train     )
					else:
						UpperCAmelCase_					,					UpperCAmelCase_     :      str						=  down_block(_snake_case			,_snake_case			,deterministic=not train     )
					down_block_res_samples += res_samples

				# 4. mid
				UpperCAmelCase_     :      str						=  self.mid_block(_snake_case			,_snake_case			,_snake_case			,deterministic=not train     )

				# 5. contronet blocks
				UpperCAmelCase_     :      int						=  ()
				for down_block_res_sample, controlnet_block in zip(_snake_case			,self.controlnet_down_blocks     ):
					UpperCAmelCase_     :      List[str]						=  controlnet_block(_snake_case     )
					controlnet_down_block_res_samples += (down_block_res_sample,)

				UpperCAmelCase_     :      str						=  controlnet_down_block_res_samples

				UpperCAmelCase_     :      List[Any]						=  self.controlnet_mid_block(_snake_case     )

				# 6. scaling
				UpperCAmelCase_     :      Any						=  [sample * conditioning_scale for sample in down_block_res_samples]
				mid_block_res_sample *= conditioning_scale

				if not return_dict:
					return (down_block_res_samples, mid_block_res_sample)

				return FlaxControlNetOutput(
				    down_block_res_samples=_snake_case			,mid_block_res_sample=_snake_case     )









import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

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


if is_torch_available():
	import torch

	from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
	from transformers.models.esm.modeling_esm import (
	    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
	    EsmEmbeddings,
	    create_position_ids_from_input_ids,
	)


class 	__snake_case   :
		def __init__(							self    ,a_    ,a_=13    ,a_=7    ,a_=False    ,a_=True    ,a_=False    ,a_=True    ,a_=33    ,a_=32    ,a_=5    ,a_=4    ,a_=37    ,a_="gelu"    ,a_=0.1    ,a_=0.1    ,a_=512    ,a_=16    ,a_=2    ,a_=0.02    ,a_=3    ,a_=4    ,a_=None    ,):







						"""simple docstring"""
						lowerCAmelCase__		  =  parent
						lowerCAmelCase__		  =  batch_size
						lowerCAmelCase__		  =  seq_length
						lowerCAmelCase__		  =  is_training
						lowerCAmelCase__		  =  use_input_mask
						lowerCAmelCase__		  =  use_token_type_ids
						lowerCAmelCase__		  =  use_labels
						lowerCAmelCase__		  =  vocab_size
						lowerCAmelCase__		  =  hidden_size
						lowerCAmelCase__		  =  num_hidden_layers
						lowerCAmelCase__		  =  num_attention_heads
						lowerCAmelCase__		  =  intermediate_size
						lowerCAmelCase__		  =  hidden_act
						lowerCAmelCase__		  =  hidden_dropout_prob
						lowerCAmelCase__		  =  attention_probs_dropout_prob
						lowerCAmelCase__		  =  max_position_embeddings
						lowerCAmelCase__		  =  type_vocab_size
						lowerCAmelCase__		  =  type_sequence_label_size
						lowerCAmelCase__		  =  initializer_range
						lowerCAmelCase__		  =  num_labels
						lowerCAmelCase__		  =  num_choices
						lowerCAmelCase__		  =  scope
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  ids_tensor([self.batch_size, self.seq_length]    ,self.vocab_size						)

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

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

						lowerCAmelCase__		  =  self.get_config()

						return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						return EsmConfig(
						    vocab_size=self.vocab_size    ,hidden_size=self.hidden_size    ,pad_token_id=1    ,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    ,initializer_range=self.initializer_range    ,)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  EsmModel(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_    ,attention_mask=a_						)
						lowerCAmelCase__		  =  model(a_						)
						lowerCAmelCase__		  =  model(a_						)

						self.parent.assertEqual(result.last_hidden_state.shape    ,(self.batch_size, self.seq_length, self.hidden_size)						)
						self.parent.assertEqual(result.pooler_output.shape    ,(self.batch_size, self.hidden_size)						)
		def 						SCREAMING_SNAKE_CASE_					(							self    ,a_    ,a_    ,a_    ,a_    ,a_    ,a_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  EsmForMaskedLM(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_    ,attention_mask=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_						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.num_labels
						lowerCAmelCase__		  =  EsmForTokenClassification(config=a_						)
						model.to(a_						)
						model.eval()
						lowerCAmelCase__		  =  model(a_    ,attention_mask=a_    ,labels=a_						)
						self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.seq_length, self.num_labels)						)



		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.prepare_config_and_inputs()
						(
						    (
						    lowerCAmelCase__
						)  ,   (
						    lowerCAmelCase__
						)  ,   (
						    lowerCAmelCase__
						)  ,   (
						    lowerCAmelCase__
						)  ,   (
						    lowerCAmelCase__
						)  ,   (
						    lowerCAmelCase__
						)  ,   
						)		  =  config_and_inputs
						lowerCAmelCase__		  =  {'input_ids': input_ids, 'attention_mask': input_mask}
						return config, inputs_dict



@require_torch
class 	__snake_case   ( SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, unittest.TestCase ):
		SCREAMING_SNAKE_CASE__								=	False

		SCREAMING_SNAKE_CASE__								=	(
		    (
		        EsmForMaskedLM,
		        EsmModel,
		        EsmForSequenceClassification,
		        EsmForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		SCREAMING_SNAKE_CASE__								=	()
		SCREAMING_SNAKE_CASE__								=	(
		    {
		        'feature-extraction': EsmModel,
		        'fill-mask': EsmForMaskedLM,
		        'text-classification': EsmForSequenceClassification,
		        'token-classification': EsmForTokenClassification,
		        'zero-shot': EsmForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		SCREAMING_SNAKE_CASE__								=	True
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  EsmModelTester(self						)
						lowerCAmelCase__		  =  ConfigTester(self    ,config_class=a_    ,hidden_size=37						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						self.config_tester.run_common_tests()
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
										lowerCAmelCase__		  =  type
										self.model_tester.create_and_check_model(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*a_						)
		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase__		  =  EsmModel.from_pretrained(a_						)
										self.assertIsNotNone(a_						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()[0]
						lowerCAmelCase__		  =  EsmEmbeddings(config=a_						)

						lowerCAmelCase__		  =  torch.as_tensor([[12, 31, 13, model.padding_idx]]						)
						lowerCAmelCase__		  =  torch.as_tensor(
						    [
						        [
						            0 + model.padding_idx + 1,
						            1 + model.padding_idx + 1,
						            2 + model.padding_idx + 1,
						            model.padding_idx,
						        ]
						    ]						)
						lowerCAmelCase__		  =  create_position_ids_from_input_ids(a_    ,model.padding_idx						)
						self.assertEqual(position_ids.shape    ,expected_positions.shape						)
						self.assertTrue(torch.all(torch.eq(a_    ,a_						)						)						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						lowerCAmelCase__		  =  self.model_tester.prepare_config_and_inputs()[0]
						lowerCAmelCase__		  =  EsmEmbeddings(config=a_						)

						lowerCAmelCase__		  =  torch.empty(2    ,4    ,30						)
						lowerCAmelCase__		  =  [
						    0 + embeddings.padding_idx + 1,
						    1 + embeddings.padding_idx + 1,
						    2 + embeddings.padding_idx + 1,
						    3 + embeddings.padding_idx + 1,
						]
						lowerCAmelCase__		  =  torch.as_tensor([expected_single_positions, expected_single_positions]						)
						lowerCAmelCase__		  =  embeddings.create_position_ids_from_inputs_embeds(a_						)
						self.assertEqual(position_ids.shape    ,expected_positions.shape						)
						self.assertTrue(torch.all(torch.eq(a_    ,a_						)						)						)
		@unittest.skip('Esm does not support embedding resizing'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass
		@unittest.skip('Esm does not support embedding resizing'						)
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						pass



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







						"""simple docstring"""
						pass



@require_torch
class 	__snake_case   ( SCREAMING_SNAKE_CASE ):
		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						with torch.no_grad():
										lowerCAmelCase__		  =  EsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'						)
										model.eval()
										lowerCAmelCase__		  =  torch.tensor([[0, 1, 2, 3, 4, 5]]						)
										lowerCAmelCase__		  =  model(a_						)[0]

										lowerCAmelCase__		  =  33

										lowerCAmelCase__		  =  torch.Size((1, 6, vocab_size)						)
										self.assertEqual(output.shape    ,a_						)

										lowerCAmelCase__		  =  torch.tensor(
										    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]						)
										self.assertTrue(torch.allclose(output[:, :3, :3]    ,a_    ,atol=1e-4						)						)



		@slow
		def 						SCREAMING_SNAKE_CASE_					(							self						):







						"""simple docstring"""
						with torch.no_grad():
										lowerCAmelCase__		  =  EsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'						)
										model.eval()

										lowerCAmelCase__		  =  torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
										lowerCAmelCase__		  =  model(a_						)[0]
										# compare the actual values for a slice.
										lowerCAmelCase__		  =  torch.tensor(
										    [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]						)
										self.assertTrue(torch.allclose(output[:, :3, :3]    ,a_    ,atol=1e-4						)						)
