





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 ):
     __A	  =       42
     __A	  =       42
     __A	  =       0.0
     __A	  =       1
     __A	  =       1
     __A	  =       True
     __A	  =       False
     __A	  =       False
     __A	  =       False
     __A	  =       jnp.floataa

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						[]
            _UpperCamelCase              =						[]

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

                   _UpperCamelCase              =						FlaxResnetBlockaD(
                       in_channels=A,       out_channels=self.out_channels,       dropout_prob=self.dropout,       dtype=self.dtype,       )
                   resnets.append(A      )

                   _UpperCamelCase              =						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(A      )

            _UpperCamelCase              =						resnets
            _UpperCamelCase              =						attentions

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



     def __call__(      self,       A,       A,       A,       A=True      ):





            """simple docstring"""

            _UpperCamelCase              =						()

            for resnet, attn in zip(self.resnets,       self.attentions      ):
                   _UpperCamelCase              =						resnet(A,       A,       deterministic=A      )
                   _UpperCamelCase              =						attn(A,       A,       deterministic=A      )
                   output_states += (hidden_states,)

            if self.add_downsample:
                   _UpperCamelCase              =						self.downsamplers_a(A      )
                   output_states += (hidden_states,)

            return hidden_states, output_states






class  __A(						nn.Module ):
     __A	  =       42
     __A	  =       42
     __A	  =       0.0
     __A	  =       1
     __A	  =       True
     __A	  =       jnp.floataa

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						[]

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

                   _UpperCamelCase              =						FlaxResnetBlockaD(
                       in_channels=A,       out_channels=self.out_channels,       dropout_prob=self.dropout,       dtype=self.dtype,       )
                   resnets.append(A      )
            _UpperCamelCase              =						resnets

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



     def __call__(      self,       A,       A,       A=True      ):





            """simple docstring"""

            _UpperCamelCase              =						()

            for resnet in self.resnets:
                   _UpperCamelCase              =						resnet(A,       A,       deterministic=A      )
                   output_states += (hidden_states,)

            if self.add_downsample:
                   _UpperCamelCase              =						self.downsamplers_a(A      )
                   output_states += (hidden_states,)

            return hidden_states, output_states






class  __A(						nn.Module ):
     __A	  =       42
     __A	  =       42
     __A	  =       42
     __A	  =       0.0
     __A	  =       1
     __A	  =       1
     __A	  =       True
     __A	  =       False
     __A	  =       False
     __A	  =       False
     __A	  =       jnp.floataa

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						[]
            _UpperCamelCase              =						[]

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

                   _UpperCamelCase              =						FlaxResnetBlockaD(
                       in_channels=resnet_in_channels + res_skip_channels,       out_channels=self.out_channels,       dropout_prob=self.dropout,       dtype=self.dtype,       )
                   resnets.append(A      )

                   _UpperCamelCase              =						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(A      )

            _UpperCamelCase              =						resnets
            _UpperCamelCase              =						attentions

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



     def __call__(      self,       A,       A,       A,       A,       A=True      ):





            """simple docstring"""

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

                   _UpperCamelCase              =						resnet(A,       A,       deterministic=A      )
                   _UpperCamelCase              =						attn(A,       A,       deterministic=A      )

            if self.add_upsample:
                   _UpperCamelCase              =						self.upsamplers_a(A      )

            return hidden_states






class  __A(						nn.Module ):
     __A	  =       42
     __A	  =       42
     __A	  =       42
     __A	  =       0.0
     __A	  =       1
     __A	  =       True
     __A	  =       jnp.floataa

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						[]

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

                   _UpperCamelCase              =						FlaxResnetBlockaD(
                       in_channels=resnet_in_channels + res_skip_channels,       out_channels=self.out_channels,       dropout_prob=self.dropout,       dtype=self.dtype,       )
                   resnets.append(A      )

            _UpperCamelCase              =						resnets

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



     def __call__(      self,       A,       A,       A,       A=True      ):





            """simple docstring"""

            for resnet in self.resnets:
                   # pop res hidden states
                   _UpperCamelCase              =						res_hidden_states_tuple[-1]
                   _UpperCamelCase              =						res_hidden_states_tuple[:-1]
                   _UpperCamelCase              =						jnp.concatenate((hidden_states, res_hidden_states),       axis=-1      )

                   _UpperCamelCase              =						resnet(A,       A,       deterministic=A      )

            if self.add_upsample:
                   _UpperCamelCase              =						self.upsamplers_a(A      )

            return hidden_states






class  __A(						nn.Module ):
     __A	  =       42
     __A	  =       0.0
     __A	  =       1
     __A	  =       1
     __A	  =       False
     __A	  =       False
     __A	  =       jnp.floataa

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						[
                FlaxResnetBlockaD(
                    in_channels=self.in_channels,       out_channels=self.in_channels,       dropout_prob=self.dropout,       dtype=self.dtype,       )
            ]

            _UpperCamelCase              =						[]

            for _ in range(self.num_layers      ):
                   _UpperCamelCase              =						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(A      )

                   _UpperCamelCase              =						FlaxResnetBlockaD(
                       in_channels=self.in_channels,       out_channels=self.in_channels,       dropout_prob=self.dropout,       dtype=self.dtype,       )
                   resnets.append(A      )

            _UpperCamelCase              =						resnets
            _UpperCamelCase              =						attentions



     def __call__(      self,       A,       A,       A,       A=True      ):





            """simple docstring"""

            _UpperCamelCase              =						self.resnets[0](A,       A      )
            for attn, resnet in zip(self.attentions,       self.resnets[1:]      ):
                   _UpperCamelCase              =						attn(A,       A,       deterministic=A      )
                   _UpperCamelCase              =						resnet(A,       A,       deterministic=A      )

            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  __A(						unittest.TestCase ):
     __A	  =       MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
     __A	  =       TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

     def 	_UpperCamelCase     (      self,       A,       A,       A      ):





            """simple docstring"""

            _UpperCamelCase              =						TextaTextGenerationPipeline(model=A,       tokenizer=A      )
            return generator, ["Something to write", "Something else"]

     def 	_UpperCamelCase     (      self,       A,       A      ):





            """simple docstring"""

            _UpperCamelCase              =						generator('''Something there'''      )
            self.assertEqual(A,       [{'''generated_text''': ANY(A      )}]      )
            # These are encoder decoder, they don't just append to incoming string
            self.assertFalse(outputs[0]['''generated_text'''].startswith('''Something there'''      )      )

            _UpperCamelCase              =						generator(['''This is great !''', '''Something else'''],       num_return_sequences=2,       do_sample=A      )
            self.assertEqual(
                A,       [
                    [{'''generated_text''': ANY(A      )}, {'''generated_text''': ANY(A      )}],
                    [{'''generated_text''': ANY(A      )}, {'''generated_text''': ANY(A      )}],
                ],       )

            _UpperCamelCase              =						generator(
                ['''This is great !''', '''Something else'''],       num_return_sequences=2,       batch_size=2,       do_sample=A      )
            self.assertEqual(
                A,       [
                    [{'''generated_text''': ANY(A      )}, {'''generated_text''': ANY(A      )}],
                    [{'''generated_text''': ANY(A      )}, {'''generated_text''': ANY(A      )}],
                ],       )

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

     @require_torch
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						pipeline('''text2text-generation''',       model='''patrickvonplaten/t5-tiny-random''',       framework='''pt'''      )
            # do_sample=False necessary for reproducibility
            _UpperCamelCase              =						generator('''Something there''',       do_sample=A      )
            self.assertEqual(A,       [{'''generated_text''': ''''''}]      )

            _UpperCamelCase              =						3
            _UpperCamelCase              =						generator(
                '''Something there''',       num_return_sequences=A,       num_beams=A,       )
            _UpperCamelCase              =						[
                {'''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(A,       A      )

            _UpperCamelCase              =						generator('''This is a test''',       do_sample=A,       num_return_sequences=2,       return_tensors=A      )
            self.assertEqual(
                A,       [
                    {'''generated_token_ids''': ANY(torch.Tensor      )},
                    {'''generated_token_ids''': ANY(torch.Tensor      )},
                ],       )
            _UpperCamelCase              =						generator.model.config.eos_token_id
            _UpperCamelCase              =						'''<pad>'''
            _UpperCamelCase              =						generator(
                ['''This is a test''', '''This is a second test'''],       do_sample=A,       num_return_sequences=2,       batch_size=2,       return_tensors=A,       )
            self.assertEqual(
                A,       [
                    [
                        {'''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 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						pipeline('''text2text-generation''',       model='''patrickvonplaten/t5-tiny-random''',       framework='''tf'''      )
            # do_sample=False necessary for reproducibility
            _UpperCamelCase              =						generator('''Something there''',       do_sample=A      )
            self.assertEqual(A,       [{'''generated_text''': ''''''}]      )






