



"""simple docstring"""






import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD



class      A__  (    nn.Module				):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      0.0
	SCREAMING_SNAKE_CASE        =      1
	SCREAMING_SNAKE_CASE        =      1
	SCREAMING_SNAKE_CASE        =      True
	SCREAMING_SNAKE_CASE        =      False
	SCREAMING_SNAKE_CASE        =      False
	SCREAMING_SNAKE_CASE        =      False
	SCREAMING_SNAKE_CASE        =      jnp.floataa

	def 				_SCREAMING_SNAKE_CASE  (	self:						Optional[int])							->			List[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    Dict				=			[]
				__lowerCAmelCase :    Optional[int]				=			[]

				for i in range(self.num_layers):
							__lowerCAmelCase :    Tuple				=			self.in_channels if i == 0 else self.out_channels

							__lowerCAmelCase :    Union[str, Any]				=			FlaxResnetBlockaD(
							    in_channels=_SCREAMING_SNAKE_CASE			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
							resnets.append(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase :    Tuple				=			FlaxTransformeraDModel(
							    in_channels=self.out_channels			,     n_heads=self.num_attention_heads			,     d_head=self.out_channels // self.num_attention_heads			,     depth=1			,     use_linear_projection=self.use_linear_projection			,     only_cross_attention=self.only_cross_attention			,     use_memory_efficient_attention=self.use_memory_efficient_attention			,     dtype=self.dtype			,     )
							attentions.append(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Union[str, Any]				=			resnets
				__lowerCAmelCase :    Dict				=			attentions

				if self.add_downsample:
							__lowerCAmelCase :    Dict				=			FlaxDownsampleaD(self.out_channels			,     dtype=self.dtype)







	def __call__(	self:						Any			,     _SCREAMING_SNAKE_CASE:						List[Any]			,     _SCREAMING_SNAKE_CASE:						List[Any]			,     _SCREAMING_SNAKE_CASE:						List[Any]			,     _SCREAMING_SNAKE_CASE:						str=True)							->			List[str]:

				"""simple docstring"""

				__lowerCAmelCase :    Optional[int]				=			()

				for resnet, attn in zip(self.resnets			,     self.attentions):
							__lowerCAmelCase :    List[Any]				=			resnet(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     deterministic=_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase :    Dict				=			attn(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     deterministic=_SCREAMING_SNAKE_CASE)
							output_states += (hidden_states,)

				if self.add_downsample:
							__lowerCAmelCase :    Dict				=			self.downsamplers_a(_SCREAMING_SNAKE_CASE)
							output_states += (hidden_states,)

				return hidden_states, output_states






class      A__  (    nn.Module				):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      0.0
	SCREAMING_SNAKE_CASE        =      1
	SCREAMING_SNAKE_CASE        =      True
	SCREAMING_SNAKE_CASE        =      jnp.floataa

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

				"""simple docstring"""

				__lowerCAmelCase :    List[Any]				=			[]

				for i in range(self.num_layers):
							__lowerCAmelCase :    str				=			self.in_channels if i == 0 else self.out_channels

							__lowerCAmelCase :    int				=			FlaxResnetBlockaD(
							    in_channels=_SCREAMING_SNAKE_CASE			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
							resnets.append(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Optional[int]				=			resnets

				if self.add_downsample:
							__lowerCAmelCase :    Optional[Any]				=			FlaxDownsampleaD(self.out_channels			,     dtype=self.dtype)







	def __call__(	self:						List[Any]			,     _SCREAMING_SNAKE_CASE:						Optional[int]			,     _SCREAMING_SNAKE_CASE:						Any			,     _SCREAMING_SNAKE_CASE:						str=True)							->			List[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    Any				=			()

				for resnet in self.resnets:
							__lowerCAmelCase :    Dict				=			resnet(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     deterministic=_SCREAMING_SNAKE_CASE)
							output_states += (hidden_states,)

				if self.add_downsample:
							__lowerCAmelCase :    Union[str, Any]				=			self.downsamplers_a(_SCREAMING_SNAKE_CASE)
							output_states += (hidden_states,)

				return hidden_states, output_states






class      A__  (    nn.Module				):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      0.0
	SCREAMING_SNAKE_CASE        =      1
	SCREAMING_SNAKE_CASE        =      1
	SCREAMING_SNAKE_CASE        =      True
	SCREAMING_SNAKE_CASE        =      False
	SCREAMING_SNAKE_CASE        =      False
	SCREAMING_SNAKE_CASE        =      False
	SCREAMING_SNAKE_CASE        =      jnp.floataa

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

				"""simple docstring"""

				__lowerCAmelCase :    List[str]				=			[]
				__lowerCAmelCase :    Any				=			[]

				for i in range(self.num_layers):
							__lowerCAmelCase :    str				=			self.in_channels if (i == self.num_layers - 1) else self.out_channels
							__lowerCAmelCase :    Union[str, Any]				=			self.prev_output_channel if i == 0 else self.out_channels

							__lowerCAmelCase :    Dict				=			FlaxResnetBlockaD(
							    in_channels=resnet_in_channels + res_skip_channels			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
							resnets.append(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase :    Union[str, Any]				=			FlaxTransformeraDModel(
							    in_channels=self.out_channels			,     n_heads=self.num_attention_heads			,     d_head=self.out_channels // self.num_attention_heads			,     depth=1			,     use_linear_projection=self.use_linear_projection			,     only_cross_attention=self.only_cross_attention			,     use_memory_efficient_attention=self.use_memory_efficient_attention			,     dtype=self.dtype			,     )
							attentions.append(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    List[str]				=			resnets
				__lowerCAmelCase :    Any				=			attentions

				if self.add_upsample:
							__lowerCAmelCase :    Union[str, Any]				=			FlaxUpsampleaD(self.out_channels			,     dtype=self.dtype)







	def __call__(	self:						Any			,     _SCREAMING_SNAKE_CASE:						Tuple			,     _SCREAMING_SNAKE_CASE:						List[Any]			,     _SCREAMING_SNAKE_CASE:						Any			,     _SCREAMING_SNAKE_CASE:						List[str]			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]=True)							->			str:

				"""simple docstring"""

				for resnet, attn in zip(self.resnets			,     self.attentions):
							# pop res hidden states
							__lowerCAmelCase :    Tuple				=			res_hidden_states_tuple[-1]
							__lowerCAmelCase :    Union[str, Any]				=			res_hidden_states_tuple[:-1]
							__lowerCAmelCase :    Optional[Any]				=			jnp.concatenate((hidden_states, res_hidden_states)			,     axis=-1)

							__lowerCAmelCase :    List[str]				=			resnet(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     deterministic=_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase :    Dict				=			attn(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     deterministic=_SCREAMING_SNAKE_CASE)

				if self.add_upsample:
							__lowerCAmelCase :    int				=			self.upsamplers_a(_SCREAMING_SNAKE_CASE)

				return hidden_states






class      A__  (    nn.Module				):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      0.0
	SCREAMING_SNAKE_CASE        =      1
	SCREAMING_SNAKE_CASE        =      True
	SCREAMING_SNAKE_CASE        =      jnp.floataa

	def 				_SCREAMING_SNAKE_CASE  (	self:						Any)							->			Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase :    str				=			[]

				for i in range(self.num_layers):
							__lowerCAmelCase :    Tuple				=			self.in_channels if (i == self.num_layers - 1) else self.out_channels
							__lowerCAmelCase :    Dict				=			self.prev_output_channel if i == 0 else self.out_channels

							__lowerCAmelCase :    Tuple				=			FlaxResnetBlockaD(
							    in_channels=resnet_in_channels + res_skip_channels			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
							resnets.append(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Tuple				=			resnets

				if self.add_upsample:
							__lowerCAmelCase :    Optional[int]				=			FlaxUpsampleaD(self.out_channels			,     dtype=self.dtype)







	def __call__(	self:						int			,     _SCREAMING_SNAKE_CASE:						Dict			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]			,     _SCREAMING_SNAKE_CASE:						List[str]			,     _SCREAMING_SNAKE_CASE:						str=True)							->			Optional[Any]:

				"""simple docstring"""

				for resnet in self.resnets:
							# pop res hidden states
							__lowerCAmelCase :    Dict				=			res_hidden_states_tuple[-1]
							__lowerCAmelCase :    Optional[Any]				=			res_hidden_states_tuple[:-1]
							__lowerCAmelCase :    Union[str, Any]				=			jnp.concatenate((hidden_states, res_hidden_states)			,     axis=-1)

							__lowerCAmelCase :    int				=			resnet(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     deterministic=_SCREAMING_SNAKE_CASE)

				if self.add_upsample:
							__lowerCAmelCase :    List[Any]				=			self.upsamplers_a(_SCREAMING_SNAKE_CASE)

				return hidden_states






class      A__  (    nn.Module				):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE        =      4_2
	SCREAMING_SNAKE_CASE        =      0.0
	SCREAMING_SNAKE_CASE        =      1
	SCREAMING_SNAKE_CASE        =      1
	SCREAMING_SNAKE_CASE        =      False
	SCREAMING_SNAKE_CASE        =      False
	SCREAMING_SNAKE_CASE        =      jnp.floataa

	def 				_SCREAMING_SNAKE_CASE  (	self:						int)							->			Dict:

				"""simple docstring"""

				__lowerCAmelCase :    Union[str, Any]				=			[
				    FlaxResnetBlockaD(
				        in_channels=self.in_channels			,     out_channels=self.in_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
				]

				__lowerCAmelCase :    int				=			[]

				for _ in range(self.num_layers):
							__lowerCAmelCase :    str				=			FlaxTransformeraDModel(
							    in_channels=self.in_channels			,     n_heads=self.num_attention_heads			,     d_head=self.in_channels // self.num_attention_heads			,     depth=1			,     use_linear_projection=self.use_linear_projection			,     use_memory_efficient_attention=self.use_memory_efficient_attention			,     dtype=self.dtype			,     )
							attentions.append(_SCREAMING_SNAKE_CASE)

							__lowerCAmelCase :    Dict				=			FlaxResnetBlockaD(
							    in_channels=self.in_channels			,     out_channels=self.in_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
							resnets.append(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase :    Any				=			resnets
				__lowerCAmelCase :    List[str]				=			attentions







	def __call__(	self:						Dict			,     _SCREAMING_SNAKE_CASE:						Tuple			,     _SCREAMING_SNAKE_CASE:						List[Any]			,     _SCREAMING_SNAKE_CASE:						Optional[Any]			,     _SCREAMING_SNAKE_CASE:						List[Any]=True)							->			str:

				"""simple docstring"""

				__lowerCAmelCase :    Optional[int]				=			self.resnets[0](_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)
				for attn, resnet in zip(self.attentions			,     self.resnets[1:]):
							__lowerCAmelCase :    Optional[Any]				=			attn(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     deterministic=_SCREAMING_SNAKE_CASE)
							__lowerCAmelCase :    Union[str, Any]				=			resnet(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     deterministic=_SCREAMING_SNAKE_CASE)

				return hidden_states







import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow


if is_torch_available():
  import torch

  from transformers import XLMRobertaModel
@require_sentencepiece
@require_tokenizers
@require_torch
class 		SCREAMING_SNAKE_CASE__	(  unittest.TestCase     ):

  '''simple docstring'''







  @slow
  def  lowerCamelCase       (      self		:			List[str]							)						->	int:


         """simple docstring"""
         _UpperCAmelCase							=    XLMRobertaModel.from_pretrained("""xlm-roberta-base"""							)
         _UpperCAmelCase							=    torch.tensor([[0, 581, 1_0269, 83, 9_9942, 136, 6_0742, 23, 70, 8_0583, 1_8276, 2]]							)
         # The dog is cute and lives in the garden house

         _UpperCAmelCase							=    torch.Size((1, 12, 768)							)  # batch_size, sequence_length, embedding_vector_dim
         _UpperCAmelCase							=    torch.tensor(
             [[-0.0101, 0.1218, -0.0803, 0.0801, 0.1327, 0.0776, -0.1215, 0.2383, 0.3338, 0.3106, 0.0300, 0.0252]]							)
         #  xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.base')
         #  xlmr.eval()
         #  expected_output_values_last_dim = xlmr.extract_features(input_ids[0])[:, :, -1]
         with torch.no_grad():
                _UpperCAmelCase							=    model(lowerCamelCase							)["""last_hidden_state"""].detach()
         self.assertEqual(output.shape       ,				lowerCamelCase							)
         # compare the actual values for a slice of last dim
         self.assertTrue(torch.allclose(output[:, :, -1]       ,				lowerCamelCase       ,				atol=1E-3							)							)





  @slow
  def  lowerCamelCase       (      self		:			Optional[int]							)						->	Union[str, Any]:


         """simple docstring"""
         _UpperCAmelCase							=    XLMRobertaModel.from_pretrained("""xlm-roberta-large"""							)
         _UpperCAmelCase							=    torch.tensor([[0, 581, 1_0269, 83, 9_9942, 136, 6_0742, 23, 70, 8_0583, 1_8276, 2]]							)
         # The dog is cute and lives in the garden house

         _UpperCAmelCase							=    torch.Size((1, 12, 1024)							)  # batch_size, sequence_length, embedding_vector_dim
         _UpperCAmelCase							=    torch.tensor(
             [[-0.0699, -0.0318, 0.0705, -0.1241, 0.0999, -0.0520, 0.1004, -0.1838, -0.4704, 0.1437, 0.0821, 0.0126]]							)
         #  xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.large')
         #  xlmr.eval()
         #  expected_output_values_last_dim = xlmr.extract_features(input_ids[0])[:, :, -1]
         with torch.no_grad():
                _UpperCAmelCase							=    model(lowerCamelCase							)["""last_hidden_state"""].detach()
         self.assertEqual(output.shape       ,				lowerCamelCase							)
         # compare the actual values for a slice of last dim
         self.assertTrue(torch.allclose(output[:, :, -1]       ,				lowerCamelCase       ,				atol=1E-3							)							)