





"""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 					_SCREAMING_SNAKE_CASE(   nn.Module  ):
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      float							=  0.0
       SCREAMING_SNAKE_CASE_   :      int							=  1
       SCREAMING_SNAKE_CASE_   :      int							=  1
       SCREAMING_SNAKE_CASE_   :      bool							=  True
       SCREAMING_SNAKE_CASE_   :      bool							=  False
       SCREAMING_SNAKE_CASE_   :      bool							=  False
       SCREAMING_SNAKE_CASE_   :      bool							=  False
       SCREAMING_SNAKE_CASE_   :      jnp.dtype							=  jnp.floataa




       def     _UpperCamelCase     (						self							)			->      Dict:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Any						=  []
             __SCREAMING_SNAKE_CASE   :str						=  []

             for i in range(self.num_layers							):
                   __SCREAMING_SNAKE_CASE   :Optional[Any]						=  self.in_channels if i == 0 else self.out_channels

                   __SCREAMING_SNAKE_CASE   :int						=  FlaxResnetBlockaD(
                       in_channels=lowerCAmelCase__		,out_channels=self.out_channels		,dropout_prob=self.dropout		,dtype=self.dtype		,)
                   resnets.append(lowerCAmelCase__							)

                   __SCREAMING_SNAKE_CASE   :int						=  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(lowerCAmelCase__							)

             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  resnets
             __SCREAMING_SNAKE_CASE   :str						=  attentions

             if self.add_downsample:
                   __SCREAMING_SNAKE_CASE   :List[Any]						=  FlaxDownsampleaD(self.out_channels		,dtype=self.dtype							)




       def __call__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=True							)			->      Dict:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  ()

             for resnet, attn in zip(self.resnets		,self.attentions							):
                   __SCREAMING_SNAKE_CASE   :int						=  resnet(lowerCAmelCase__		,lowerCAmelCase__		,deterministic=lowerCAmelCase__							)
                   __SCREAMING_SNAKE_CASE   :str						=  attn(lowerCAmelCase__		,lowerCAmelCase__		,deterministic=lowerCAmelCase__							)
                   output_states += (hidden_states,)

             if self.add_downsample:
                   __SCREAMING_SNAKE_CASE   :Optional[Any]						=  self.downsamplers_a(lowerCAmelCase__							)
                   output_states += (hidden_states,)

             return hidden_states, output_states






class 					_SCREAMING_SNAKE_CASE(   nn.Module  ):
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      float							=  0.0
       SCREAMING_SNAKE_CASE_   :      int							=  1
       SCREAMING_SNAKE_CASE_   :      bool							=  True
       SCREAMING_SNAKE_CASE_   :      jnp.dtype							=  jnp.floataa




       def     _UpperCamelCase     (						self							)			->      Optional[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :List[str]						=  []

             for i in range(self.num_layers							):
                   __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.in_channels if i == 0 else self.out_channels

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

             if self.add_downsample:
                   __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  FlaxDownsampleaD(self.out_channels		,dtype=self.dtype							)




       def __call__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=True							)			->      Any:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :int						=  ()

             for resnet in self.resnets:
                   __SCREAMING_SNAKE_CASE   :Tuple						=  resnet(lowerCAmelCase__		,lowerCAmelCase__		,deterministic=lowerCAmelCase__							)
                   output_states += (hidden_states,)

             if self.add_downsample:
                   __SCREAMING_SNAKE_CASE   :Any						=  self.downsamplers_a(lowerCAmelCase__							)
                   output_states += (hidden_states,)

             return hidden_states, output_states






class 					_SCREAMING_SNAKE_CASE(   nn.Module  ):
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      float							=  0.0
       SCREAMING_SNAKE_CASE_   :      int							=  1
       SCREAMING_SNAKE_CASE_   :      int							=  1
       SCREAMING_SNAKE_CASE_   :      bool							=  True
       SCREAMING_SNAKE_CASE_   :      bool							=  False
       SCREAMING_SNAKE_CASE_   :      bool							=  False
       SCREAMING_SNAKE_CASE_   :      bool							=  False
       SCREAMING_SNAKE_CASE_   :      jnp.dtype							=  jnp.floataa




       def     _UpperCamelCase     (						self							)			->      str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :int						=  []
             __SCREAMING_SNAKE_CASE   :List[Any]						=  []

             for i in range(self.num_layers							):
                   __SCREAMING_SNAKE_CASE   :int						=  self.in_channels if (i == self.num_layers - 1) else self.out_channels
                   __SCREAMING_SNAKE_CASE   :List[str]						=  self.prev_output_channel if i == 0 else self.out_channels

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

                   __SCREAMING_SNAKE_CASE   :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(lowerCAmelCase__							)

             __SCREAMING_SNAKE_CASE   :int						=  resnets
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  attentions

             if self.add_upsample:
                   __SCREAMING_SNAKE_CASE   :str						=  FlaxUpsampleaD(self.out_channels		,dtype=self.dtype							)




       def __call__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=True							)			->      Optional[int]:





             """simple docstring"""






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

                   __SCREAMING_SNAKE_CASE   :Tuple						=  resnet(lowerCAmelCase__		,lowerCAmelCase__		,deterministic=lowerCAmelCase__							)
                   __SCREAMING_SNAKE_CASE   :List[Any]						=  attn(lowerCAmelCase__		,lowerCAmelCase__		,deterministic=lowerCAmelCase__							)

             if self.add_upsample:
                   __SCREAMING_SNAKE_CASE   :str						=  self.upsamplers_a(lowerCAmelCase__							)

             return hidden_states






class 					_SCREAMING_SNAKE_CASE(   nn.Module  ):
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      float							=  0.0
       SCREAMING_SNAKE_CASE_   :      int							=  1
       SCREAMING_SNAKE_CASE_   :      bool							=  True
       SCREAMING_SNAKE_CASE_   :      jnp.dtype							=  jnp.floataa




       def     _UpperCamelCase     (						self							)			->      List[str]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :str						=  []

             for i in range(self.num_layers							):
                   __SCREAMING_SNAKE_CASE   :int						=  self.in_channels if (i == self.num_layers - 1) else self.out_channels
                   __SCREAMING_SNAKE_CASE   :List[str]						=  self.prev_output_channel if i == 0 else self.out_channels

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

             __SCREAMING_SNAKE_CASE   :int						=  resnets

             if self.add_upsample:
                   __SCREAMING_SNAKE_CASE   :Dict						=  FlaxUpsampleaD(self.out_channels		,dtype=self.dtype							)




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





             """simple docstring"""






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

                   __SCREAMING_SNAKE_CASE   :Dict						=  resnet(lowerCAmelCase__		,lowerCAmelCase__		,deterministic=lowerCAmelCase__							)

             if self.add_upsample:
                   __SCREAMING_SNAKE_CASE   :Optional[Any]						=  self.upsamplers_a(lowerCAmelCase__							)

             return hidden_states






class 					_SCREAMING_SNAKE_CASE(   nn.Module  ):
       SCREAMING_SNAKE_CASE_   :      int
       SCREAMING_SNAKE_CASE_   :      float							=  0.0
       SCREAMING_SNAKE_CASE_   :      int							=  1
       SCREAMING_SNAKE_CASE_   :      int							=  1
       SCREAMING_SNAKE_CASE_   :      bool							=  False
       SCREAMING_SNAKE_CASE_   :      bool							=  False
       SCREAMING_SNAKE_CASE_   :      jnp.dtype							=  jnp.floataa




       def     _UpperCamelCase     (						self							)			->      List[Any]:





             """simple docstring"""






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

             __SCREAMING_SNAKE_CASE   :List[Any]						=  []

             for _ in range(self.num_layers							):
                   __SCREAMING_SNAKE_CASE   :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(lowerCAmelCase__							)

                   __SCREAMING_SNAKE_CASE   :Any						=  FlaxResnetBlockaD(
                       in_channels=self.in_channels		,out_channels=self.in_channels		,dropout_prob=self.dropout		,dtype=self.dtype		,)
                   resnets.append(lowerCAmelCase__							)

             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  resnets
             __SCREAMING_SNAKE_CASE   :Dict						=  attentions




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.resnets[0](lowerCAmelCase__		,lowerCAmelCase__							)
             for attn, resnet in zip(self.attentions		,self.resnets[1:]							):
                   __SCREAMING_SNAKE_CASE   :Any						=  attn(lowerCAmelCase__		,lowerCAmelCase__		,deterministic=lowerCAmelCase__							)
                   __SCREAMING_SNAKE_CASE   :Any						=  resnet(lowerCAmelCase__		,lowerCAmelCase__		,deterministic=lowerCAmelCase__							)

             return hidden_states


import unittest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TextaTextGenerationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, require_tf, require_torch
from transformers.utils import is_torch_available

from .test_pipelines_common import ANY


if is_torch_available():
						import torch



@is_pipeline_test
class   __lowercase				( unittest.TestCase	):




			"""simple docstring"""







			_UpperCAmelCase   :							List[Any]										=     MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
			_UpperCAmelCase   :							str										=     TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING



			def        _SCREAMING_SNAKE_CASE			( self  :     Any   ,						lowerCAmelCase__  :     List[str]   ,						lowerCAmelCase__  :     List[str]   ,						lowerCAmelCase__  :     Dict):
									SCREAMING_SNAKE_CASE_: Any  		=       TextaTextGenerationPipeline(model=lowerCAmelCase__   ,						tokenizer=lowerCAmelCase__)
									return generator, ["Something to write", "Something else"]



			def        _SCREAMING_SNAKE_CASE			( self  :     str   ,						lowerCAmelCase__  :     Dict   ,						lowerCAmelCase__  :     Any):
									SCREAMING_SNAKE_CASE_: List[Any]  		=       generator("Something there")
									self.assertEqual(lowerCAmelCase__   ,						[{"generated_text": ANY(lowerCAmelCase__)}])
									# These are encoder decoder, they don't just append to incoming string
									self.assertFalse(outputs[0]["generated_text"].startswith("Something there"))

									SCREAMING_SNAKE_CASE_: List[Any]  		=       generator(["This is great !", "Something else"]   ,						num_return_sequences=2   ,						do_sample=lowerCAmelCase__)
									self.assertEqual(
									    lowerCAmelCase__   ,						[
									        [{"generated_text": ANY(lowerCAmelCase__)}, {"generated_text": ANY(lowerCAmelCase__)}],
									        [{"generated_text": ANY(lowerCAmelCase__)}, {"generated_text": ANY(lowerCAmelCase__)}],
									    ]   ,						)

									SCREAMING_SNAKE_CASE_: Dict  		=       generator(
									    ["This is great !", "Something else"]   ,						num_return_sequences=2   ,						batch_size=2   ,						do_sample=lowerCAmelCase__)
									self.assertEqual(
									    lowerCAmelCase__   ,						[
									        [{"generated_text": ANY(lowerCAmelCase__)}, {"generated_text": ANY(lowerCAmelCase__)}],
									        [{"generated_text": ANY(lowerCAmelCase__)}, {"generated_text": ANY(lowerCAmelCase__)}],
									    ]   ,						)

									with self.assertRaises(lowerCAmelCase__):
															generator(4)



			@require_torch
			def        _SCREAMING_SNAKE_CASE			( self  :     Tuple):
									SCREAMING_SNAKE_CASE_: Optional[int]  		=       pipeline("text2text-generation"   ,						model="patrickvonplaten/t5-tiny-random"   ,						framework="pt")
									# do_sample=False necessary for reproducibility
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       generator("Something there"   ,						do_sample=lowerCAmelCase__)
									self.assertEqual(lowerCAmelCase__   ,						[{"generated_text": ""}])

									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       3
									SCREAMING_SNAKE_CASE_: Any  		=       generator(
									    "Something there"   ,						num_return_sequences=lowerCAmelCase__   ,						num_beams=lowerCAmelCase__   ,						)
									SCREAMING_SNAKE_CASE_: Any  		=       [
									    {"generated_text": "Beide Beide Beide Beide Beide Beide Beide Beide Beide"},
									    {"generated_text": "Beide Beide Beide Beide Beide Beide Beide Beide"},
									    {"generated_text": ""},
									]
									self.assertEqual(lowerCAmelCase__   ,						lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_: int  		=       generator("This is a test"   ,						do_sample=lowerCAmelCase__   ,						num_return_sequences=2   ,						return_tensors=lowerCAmelCase__)
									self.assertEqual(
									    lowerCAmelCase__   ,						[
									        {"generated_token_ids": ANY(torch.Tensor)},
									        {"generated_token_ids": ANY(torch.Tensor)},
									    ]   ,						)
									SCREAMING_SNAKE_CASE_: str  		=       generator.model.config.eos_token_id
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       "<pad>"
									SCREAMING_SNAKE_CASE_: Tuple  		=       generator(
									    ["This is a test", "This is a second test"]   ,						do_sample=lowerCAmelCase__   ,						num_return_sequences=2   ,						batch_size=2   ,						return_tensors=lowerCAmelCase__   ,						)
									self.assertEqual(
									    lowerCAmelCase__   ,						[
									        [
									            {"generated_token_ids": ANY(torch.Tensor)},
									            {"generated_token_ids": ANY(torch.Tensor)},
									        ],
									        [
									            {"generated_token_ids": ANY(torch.Tensor)},
									            {"generated_token_ids": ANY(torch.Tensor)},
									        ],
									    ]   ,						)




			@require_tf
			def        _SCREAMING_SNAKE_CASE			( self  :     Union[str, Any]):
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       pipeline("text2text-generation"   ,						model="patrickvonplaten/t5-tiny-random"   ,						framework="tf")
									# do_sample=False necessary for reproducibility
									SCREAMING_SNAKE_CASE_: List[Any]  		=       generator("Something there"   ,						do_sample=lowerCAmelCase__)
									self.assertEqual(lowerCAmelCase__   ,						[{"generated_text": ""}])






