





'''simple docstring'''




from __future__ import annotations

import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import is_tf_available, is_vision_available

from ...test_modeling_tf_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_tf_bert import TFBertModelTester
from ..clip.test_modeling_tf_clip import TFCLIPVisionModelTester
from ..deit.test_modeling_tf_deit import TFDeiTModelTester
from ..roberta.test_modeling_tf_roberta import TFRobertaModelTester
from ..vit.test_modeling_tf_vit import TFViTModelTester


if is_tf_available():
  from transformers import (
      TFBertModel,
      TFCLIPVisionModel,
      TFDeiTModel,
      TFRobertaModel,
      TFVisionTextDualEncoderModel,
      TFViTModel,
      VisionTextDualEncoderConfig,
  )

if is_vision_available():
  from PIL import Image

  from transformers import VisionTextDualEncoderProcessor
def 	UpperCAmelCase_							(					__lowercase							:      Optional[Any]       )		->       Optional[Any]:


     '''simple docstring'''



     if isinstance(__lowercase			,     collections.abc.Iterable       ):
          return x
     return (x, x)







@require_tf
class   A_	:




 def      lowercase	(					self       :    str	,			snake_case_       :    int	,			snake_case_       :    int    ):
      pass




 def      lowercase	(					self       :    Optional[Any]    ):
      pass




 def      lowercase	(					self       :    int    ):
      pass




 def      lowercase	(					self       :    List[Any]	,			snake_case_       :    int	,			snake_case_       :    Dict	,			snake_case_       :    Any	,			snake_case_       :    int	,			snake_case_       :    List[str]=None	,			**snake_case_       :    int    ):
      _UpperCAmelCase     							=   VisionTextDualEncoderConfig.from_vision_text_configs(snake_case_	,			snake_case_    )

      _UpperCAmelCase     							=   TFVisionTextDualEncoderModel(snake_case_    )

      _UpperCAmelCase     							=   model(input_ids=snake_case_	,			pixel_values=snake_case_	,			attention_mask=snake_case_    )

      self.assertEqual(output["text_embeds"].shape	,			(input_ids.shape[0], config.projection_dim)    )
      self.assertEqual(output["image_embeds"].shape	,			(pixel_values.shape[0], config.projection_dim)    )




 def      lowercase	(					self       :    List[str]	,			snake_case_       :    int	,			snake_case_       :    Tuple	,			snake_case_       :    Union[str, Any]	,			snake_case_       :    Optional[int]	,			snake_case_       :    Tuple=None	,			**snake_case_       :    Optional[int]    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.get_vision_text_model(snake_case_	,			snake_case_    )
      _UpperCAmelCase     							=   TFVisionTextDualEncoderModel(vision_model=snake_case_	,			text_model=snake_case_    )

      _UpperCAmelCase     							=   model(input_ids=snake_case_	,			pixel_values=snake_case_	,			attention_mask=snake_case_    )

      self.assertEqual(output["text_embeds"].shape	,			(input_ids.shape[0], model.config.projection_dim)    )
      self.assertEqual(output["image_embeds"].shape	,			(pixel_values.shape[0], model.config.projection_dim)    )




 def      lowercase	(					self       :    Union[str, Any]	,			snake_case_       :    Union[str, Any]	,			snake_case_       :    int	,			snake_case_       :    Tuple	,			snake_case_       :    int	,			snake_case_       :    Union[str, Any]=None	,			**snake_case_       :    Any    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.get_vision_text_model(snake_case_	,			snake_case_    )
      _UpperCAmelCase     							=   {"vision_model": vision_model, "text_model": text_model}
      _UpperCAmelCase     							=   TFVisionTextDualEncoderModel.from_vision_text_pretrained(**snake_case_    )

      _UpperCAmelCase     							=   model(input_ids=snake_case_	,			pixel_values=snake_case_	,			attention_mask=snake_case_    )

      self.assertEqual(output["text_embeds"].shape	,			(input_ids.shape[0], model.config.projection_dim)    )
      self.assertEqual(output["image_embeds"].shape	,			(pixel_values.shape[0], model.config.projection_dim)    )




 def      lowercase	(					self       :    Optional[int]	,			snake_case_       :    List[Any]	,			snake_case_       :    Optional[Any]	,			snake_case_       :    Optional[Any]	,			snake_case_       :    Tuple	,			snake_case_       :    Optional[Any]=None	,			**snake_case_       :    Union[str, Any]    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.get_vision_text_model(snake_case_	,			snake_case_    )
      _UpperCAmelCase     							=   TFVisionTextDualEncoderModel(vision_model=snake_case_	,			text_model=snake_case_    )

      _UpperCAmelCase     							=   model(input_ids=snake_case_	,			pixel_values=snake_case_	,			attention_mask=snake_case_    )
      _UpperCAmelCase     							=   output[0].numpy()

      with tempfile.TemporaryDirectory() as tmpdirname:
           model.save_pretrained(snake_case_    )
           _UpperCAmelCase     							=   TFVisionTextDualEncoderModel.from_pretrained(snake_case_    )

           _UpperCAmelCase     							=   model(input_ids=snake_case_	,			pixel_values=snake_case_	,			attention_mask=snake_case_    )
           _UpperCAmelCase     							=   after_output[0].numpy()
           _UpperCAmelCase     							=   np.amax(np.abs(out_a - out_a    )    )
           self.assertLessEqual(snake_case_	,			1e-5    )




 def      lowercase	(					self       :    int	,			snake_case_       :    Optional[Any]	,			snake_case_       :    Any	,			snake_case_       :    List[Any]	,			snake_case_       :    Any	,			snake_case_       :    int=None	,			**snake_case_       :    str    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.get_vision_text_model(snake_case_	,			snake_case_    )
      _UpperCAmelCase     							=   TFVisionTextDualEncoderModel(vision_model=snake_case_	,			text_model=snake_case_    )

      _UpperCAmelCase     							=   model(
          input_ids=snake_case_	,			pixel_values=snake_case_	,			attention_mask=snake_case_	,			output_attentions=snake_case_    )

      _UpperCAmelCase     							=   output.vision_model_output.attentions
      self.assertEqual(len(snake_case_    )	,			vision_config.num_hidden_layers    )

      # in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
      _UpperCAmelCase     							=   to_atuple(vision_model.config.image_size    )
      _UpperCAmelCase     							=   to_atuple(vision_model.config.patch_size    )
      _UpperCAmelCase     							=   (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
      _UpperCAmelCase     							=   num_patches + 1
      self.assertEqual(vision_attentions[0].shape[-3:]	,			(vision_config.num_attention_heads, seq_len, seq_len)    )

      _UpperCAmelCase     							=   output.text_model_output.attentions
      self.assertEqual(len(snake_case_    )	,			text_config.num_hidden_layers    )

      self.assertEqual(
          text_attentions[0].shape[-3:]	,			(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])	,			)




 def      lowercase	(					self       :    Optional[int]	,			snake_case_       :    np.ndarray	,			snake_case_       :    np.ndarray	,			snake_case_       :    float    ):
      _UpperCAmelCase     							=   np.abs((a - b)    ).max()
      self.assertLessEqual(snake_case_	,			snake_case_	,			f'Difference between torch and flax is {diff} (>= {tol}).'    )




 def      lowercase	(					self       :    Tuple    ):
      _UpperCAmelCase     							=   self.prepare_config_and_inputs()
      self.check_vision_text_dual_encoder_model(**snake_case_    )




 def      lowercase	(					self       :    Any    ):
      _UpperCAmelCase     							=   self.prepare_config_and_inputs()
      self.check_model_from_pretrained_configs(**snake_case_    )




 def      lowercase	(					self       :    Any    ):
      _UpperCAmelCase     							=   self.prepare_config_and_inputs()
      self.check_vision_text_dual_encoder_from_pretrained(**snake_case_    )




 def      lowercase	(					self       :    Union[str, Any]    ):
      _UpperCAmelCase     							=   self.prepare_config_and_inputs()
      self.check_save_load(**snake_case_    )




 def      lowercase	(					self       :    Union[str, Any]    ):
      _UpperCAmelCase     							=   self.prepare_config_and_inputs()
      self.check_vision_text_output_attention(**snake_case_    )




 @slow
 def      lowercase	(					self       :    List[str]    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.get_pretrained_model_and_inputs()

      _UpperCAmelCase     							=   model_a(**snake_case_    )
      _UpperCAmelCase     							=   outputs[0].numpy()

      with tempfile.TemporaryDirectory() as tmp_dirname:
           model_a.save_pretrained(snake_case_    )
           _UpperCAmelCase     							=   TFVisionTextDualEncoderModel.from_pretrained(snake_case_    )

           _UpperCAmelCase     							=   model_a(**snake_case_    )
           _UpperCAmelCase     							=   after_outputs[0].numpy()
           _UpperCAmelCase     							=   np.amax(np.abs(out_a - out_a    )    )
           self.assertLessEqual(snake_case_	,			1e-5    )







@require_tf
class   A_	(  lowerCAmelCase_							,						unittest.TestCase			):




 def      lowercase	(					self       :    Optional[int]    ):
      _UpperCAmelCase     							=   TFVisionTextDualEncoderModel.from_vision_text_pretrained(
          "hf-internal-testing/tiny-random-vit"	,			"hf-internal-testing/tiny-random-bert"    )
      _UpperCAmelCase     							=   1_3
      _UpperCAmelCase     							=   floats_tensor(
          [
              batch_size,
              model.vision_model.config.num_channels,
              model.vision_model.config.image_size,
              model.vision_model.config.image_size,
          ]    )
      _UpperCAmelCase     							=   ids_tensor([batch_size, 4]	,			model.text_model.config.vocab_size    )
      _UpperCAmelCase     							=   random_attention_mask([batch_size, 4]    )
      _UpperCAmelCase     							=   {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

      return model, inputs




 def      lowercase	(					self       :    str	,			snake_case_       :    Optional[int]	,			snake_case_       :    Union[str, Any]    ):
      _UpperCAmelCase     							=   TFViTModel(snake_case_	,			name="vision_model"    )
      _UpperCAmelCase     							=   TFBertModel(snake_case_	,			name="text_model"    )
      return vision_model, text_model




 def      lowercase	(					self       :    int    ):
      _UpperCAmelCase     							=   TFViTModelTester(self    )
      _UpperCAmelCase     							=   TFBertModelTester(self    )
      _UpperCAmelCase     							=   vit_model_tester.prepare_config_and_inputs()
      _UpperCAmelCase     							=   bert_model_tester.prepare_config_and_inputs()

      _UpperCAmelCase							,  _UpperCAmelCase							,  _UpperCAmelCase     							=   vision_config_and_inputs

      (
          (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  
      )     							=   text_config_and_inputs

      return {
          "text_config": text_config,
          "vision_config": vision_config,
          "pixel_values": pixel_values,
          "attention_mask": input_mask,
          "input_ids": input_ids,
          "text_token_type_ids": token_type_ids,
          "text_sequence_labels": sequence_labels,
          "text_token_labels": token_labels,
          "text_choice_labels": choice_labels,
      }







@require_tf
class   A_	(  lowerCAmelCase_							,						unittest.TestCase			):




 def      lowercase	(					self       :    List[Any]    ):
      # DeiT repo doesn't have TF weights, but we don't actually use the weights at all so let's
      # just reinitialize it.
      _UpperCAmelCase     							=   TFVisionTextDualEncoderModel.from_vision_text_pretrained(
          "Rocketknight1/tiny-random-deit-tf"	,			"hf-internal-testing/tiny-random-roberta"    )
      _UpperCAmelCase     							=   1_3
      _UpperCAmelCase     							=   floats_tensor(
          [
              batch_size,
              model.vision_model.config.num_channels,
              model.vision_model.config.image_size,
              model.vision_model.config.image_size,
          ]    )
      _UpperCAmelCase     							=   ids_tensor([batch_size, 4]	,			model.text_model.config.vocab_size    )
      _UpperCAmelCase     							=   random_attention_mask([batch_size, 4]    )
      _UpperCAmelCase     							=   {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

      return model, inputs




 def      lowercase	(					self       :    int	,			snake_case_       :    int	,			snake_case_       :    Dict	,			snake_case_       :    str	,			snake_case_       :    Tuple	,			snake_case_       :    str=None	,			**snake_case_       :    Tuple    ):
      _UpperCAmelCase							,  _UpperCAmelCase     							=   self.get_vision_text_model(snake_case_	,			snake_case_    )
      _UpperCAmelCase     							=   TFVisionTextDualEncoderModel(vision_model=snake_case_	,			text_model=snake_case_    )

      _UpperCAmelCase     							=   model(
          input_ids=snake_case_	,			pixel_values=snake_case_	,			attention_mask=snake_case_	,			output_attentions=snake_case_    )

      _UpperCAmelCase     							=   output.vision_model_output.attentions
      self.assertEqual(len(snake_case_    )	,			vision_config.num_hidden_layers    )

      # in DEiT, the seq_len equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
      _UpperCAmelCase     							=   to_atuple(vision_model.config.image_size    )
      _UpperCAmelCase     							=   to_atuple(vision_model.config.patch_size    )
      _UpperCAmelCase     							=   (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
      _UpperCAmelCase     							=   num_patches + 2
      self.assertEqual(vision_attentions[0].shape[-3:]	,			(vision_config.num_attention_heads, seq_len, seq_len)    )

      _UpperCAmelCase     							=   output.text_model_output.attentions
      self.assertEqual(len(snake_case_    )	,			text_config.num_hidden_layers    )

      self.assertEqual(
          text_attentions[0].shape[-3:]	,			(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])	,			)




 def      lowercase	(					self       :    List[Any]	,			snake_case_       :    Dict	,			snake_case_       :    Optional[Any]    ):
      _UpperCAmelCase     							=   TFDeiTModel(snake_case_	,			name="vision_model"    )
      _UpperCAmelCase     							=   TFRobertaModel(snake_case_	,			name="text_model"    )
      return vision_model, text_model




 def      lowercase	(					self       :    Any    ):
      _UpperCAmelCase     							=   TFDeiTModelTester(self    )
      _UpperCAmelCase     							=   TFRobertaModelTester(self    )
      _UpperCAmelCase     							=   vit_model_tester.prepare_config_and_inputs()
      _UpperCAmelCase     							=   bert_model_tester.prepare_config_and_inputs()

      _UpperCAmelCase							,  _UpperCAmelCase							,  _UpperCAmelCase     							=   vision_config_and_inputs

      (
          (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  
      )     							=   text_config_and_inputs

      return {
          "text_config": text_config,
          "vision_config": vision_config,
          "pixel_values": pixel_values,
          "attention_mask": input_mask,
          "input_ids": input_ids,
          "text_token_type_ids": token_type_ids,
          "text_sequence_labels": sequence_labels,
          "text_token_labels": token_labels,
          "text_choice_labels": choice_labels,
      }







@require_tf
class   A_	(  lowerCAmelCase_							,						unittest.TestCase			):




 def      lowercase	(					self       :    Union[str, Any]    ):
      _UpperCAmelCase     							=   TFVisionTextDualEncoderModel.from_vision_text_pretrained(
          "Rocketknight1/tiny-random-clip-tf"	,			"hf-internal-testing/tiny-random-bert"    )
      _UpperCAmelCase     							=   1_3
      _UpperCAmelCase     							=   floats_tensor(
          [
              batch_size,
              model.vision_model.config.num_channels,
              model.vision_model.config.image_size,
              model.vision_model.config.image_size,
          ]    )
      _UpperCAmelCase     							=   ids_tensor([batch_size, 4]	,			model.text_model.config.vocab_size    )
      _UpperCAmelCase     							=   random_attention_mask([batch_size, 4]    )
      _UpperCAmelCase     							=   {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

      return model, inputs




 def      lowercase	(					self       :    Union[str, Any]	,			snake_case_       :    str	,			snake_case_       :    List[Any]    ):
      _UpperCAmelCase     							=   TFCLIPVisionModel(snake_case_	,			name="vision_model"    )
      _UpperCAmelCase     							=   TFBertModel(snake_case_	,			name="text_model"    )
      return vision_model, text_model




 def      lowercase	(					self       :    Optional[int]    ):
      _UpperCAmelCase     							=   TFCLIPVisionModelTester(self    )
      _UpperCAmelCase     							=   TFBertModelTester(self    )
      _UpperCAmelCase     							=   clip_model_tester.prepare_config_and_inputs()
      _UpperCAmelCase     							=   bert_model_tester.prepare_config_and_inputs()

      _UpperCAmelCase							,  _UpperCAmelCase     							=   vision_config_and_inputs

      (
          (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  (
          _UpperCAmelCase
      )							,  
      )     							=   text_config_and_inputs

      return {
          "text_config": text_config,
          "vision_config": vision_config,
          "pixel_values": pixel_values,
          "attention_mask": input_mask,
          "input_ids": input_ids,
          "text_token_type_ids": token_type_ids,
          "text_sequence_labels": sequence_labels,
          "text_token_labels": token_labels,
          "text_choice_labels": choice_labels,
      }







@require_vision
@require_tf
class   A_	(  unittest.TestCase			):




 @slow
 def      lowercase	(					self       :    int    ):
      _UpperCAmelCase     							=   TFVisionTextDualEncoderModel.from_pretrained(
          "clip-italian/clip-italian"	,			logit_scale_init_value=1.0	,			from_pt=snake_case_    )
      _UpperCAmelCase     							=   VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian"    )

      _UpperCAmelCase     							=   Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
      _UpperCAmelCase     							=   processor(
          text=["una foto di un gatto", "una foto di un cane"]	,			images=snake_case_	,			padding=snake_case_	,			return_tensors="np"    )

      _UpperCAmelCase     							=   model(**snake_case_    )

      # verify the logits
      self.assertEqual(outputs.logits_per_image.shape	,			(inputs.pixel_values.shape[0], inputs.input_ids.shape[0])    )
      self.assertEqual(
          outputs.logits_per_text.shape	,			(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])	,			)

      _UpperCAmelCase     							=   np.array([[1.2_2_8_4_7_2_7, 0.3_1_0_4_1_2_2]]    )

      self.assertTrue(np.allclose(outputs.logits_per_image.numpy()	,			snake_case_	,			atol=1e-3    )    )






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

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




class 						lowercase__  (						nn.Module  ):
	A__    :   int
	A__    :   int
	A__    :   float				       =0.0
	A__    :   int				       =1
	A__    :   int				       =1
	A__    :   bool				       =True
	A__    :   bool				       =False
	A__    :   bool				       =False
	A__    :   bool				       =False
	A__    :   jnp.dtype				       =jnp.floataa
	def      A_				(			self							:   Optional[Any]							):
						SCREAMING_SNAKE_CASE__													= []
						SCREAMING_SNAKE_CASE__													= []

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

											SCREAMING_SNAKE_CASE__													= FlaxResnetBlockaD(
											    in_channels=UpperCAmelCase_       ,   out_channels=self.out_channels       ,   dropout_prob=self.dropout       ,   dtype=self.dtype       ,   )
											resnets.append(UpperCAmelCase_							)

											SCREAMING_SNAKE_CASE__													= 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(UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= resnets
						SCREAMING_SNAKE_CASE__													= attentions

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





	def __call__(			self							:   int       ,   UpperCAmelCase_							:   str       ,   UpperCAmelCase_							:   Tuple       ,   UpperCAmelCase_							:   Dict       ,   UpperCAmelCase_							:   int=True							):
						SCREAMING_SNAKE_CASE__													= ()

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

						if self.add_downsample:
											SCREAMING_SNAKE_CASE__													= self.downsamplers_a(UpperCAmelCase_							)
											output_states += (hidden_states,)

						return hidden_states, output_states





class 						lowercase__  (						nn.Module  ):
	A__    :   int
	A__    :   int
	A__    :   float				       =0.0
	A__    :   int				       =1
	A__    :   bool				       =True
	A__    :   jnp.dtype				       =jnp.floataa
	def      A_				(			self							:   List[str]							):
						SCREAMING_SNAKE_CASE__													= []

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

											SCREAMING_SNAKE_CASE__													= FlaxResnetBlockaD(
											    in_channels=UpperCAmelCase_       ,   out_channels=self.out_channels       ,   dropout_prob=self.dropout       ,   dtype=self.dtype       ,   )
											resnets.append(UpperCAmelCase_							)
						SCREAMING_SNAKE_CASE__													= resnets

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





	def __call__(			self							:   List[str]       ,   UpperCAmelCase_							:   Any       ,   UpperCAmelCase_							:   Union[str, Any]       ,   UpperCAmelCase_							:   List[Any]=True							):
						SCREAMING_SNAKE_CASE__													= ()

						for resnet in self.resnets:
											SCREAMING_SNAKE_CASE__													= resnet(UpperCAmelCase_       ,   UpperCAmelCase_       ,   deterministic=UpperCAmelCase_							)
											output_states += (hidden_states,)

						if self.add_downsample:
											SCREAMING_SNAKE_CASE__													= self.downsamplers_a(UpperCAmelCase_							)
											output_states += (hidden_states,)

						return hidden_states, output_states





class 						lowercase__  (						nn.Module  ):
	A__    :   int
	A__    :   int
	A__    :   int
	A__    :   float				       =0.0
	A__    :   int				       =1
	A__    :   int				       =1
	A__    :   bool				       =True
	A__    :   bool				       =False
	A__    :   bool				       =False
	A__    :   bool				       =False
	A__    :   jnp.dtype				       =jnp.floataa
	def      A_				(			self							:   List[str]							):
						SCREAMING_SNAKE_CASE__													= []
						SCREAMING_SNAKE_CASE__													= []

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

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

											SCREAMING_SNAKE_CASE__													= 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(UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= resnets
						SCREAMING_SNAKE_CASE__													= attentions

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





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

											SCREAMING_SNAKE_CASE__													= resnet(UpperCAmelCase_       ,   UpperCAmelCase_       ,   deterministic=UpperCAmelCase_							)
											SCREAMING_SNAKE_CASE__													= attn(UpperCAmelCase_       ,   UpperCAmelCase_       ,   deterministic=UpperCAmelCase_							)

						if self.add_upsample:
											SCREAMING_SNAKE_CASE__													= self.upsamplers_a(UpperCAmelCase_							)

						return hidden_states





class 						lowercase__  (						nn.Module  ):
	A__    :   int
	A__    :   int
	A__    :   int
	A__    :   float				       =0.0
	A__    :   int				       =1
	A__    :   bool				       =True
	A__    :   jnp.dtype				       =jnp.floataa
	def      A_				(			self							:   List[Any]							):
						SCREAMING_SNAKE_CASE__													= []

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

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

						SCREAMING_SNAKE_CASE__													= resnets

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





	def __call__(			self							:   Tuple       ,   UpperCAmelCase_							:   Any       ,   UpperCAmelCase_							:   Any       ,   UpperCAmelCase_							:   Dict       ,   UpperCAmelCase_							:   Dict=True							):
						for resnet in self.resnets:
											# pop res hidden states
											SCREAMING_SNAKE_CASE__													= res_hidden_states_tuple[-1]
											SCREAMING_SNAKE_CASE__													= res_hidden_states_tuple[:-1]
											SCREAMING_SNAKE_CASE__													= jnp.concatenate((hidden_states, res_hidden_states)       ,   axis=-1							)

											SCREAMING_SNAKE_CASE__													= resnet(UpperCAmelCase_       ,   UpperCAmelCase_       ,   deterministic=UpperCAmelCase_							)

						if self.add_upsample:
											SCREAMING_SNAKE_CASE__													= self.upsamplers_a(UpperCAmelCase_							)

						return hidden_states





class 						lowercase__  (						nn.Module  ):
	A__    :   int
	A__    :   float				       =0.0
	A__    :   int				       =1
	A__    :   int				       =1
	A__    :   bool				       =False
	A__    :   bool				       =False
	A__    :   jnp.dtype				       =jnp.floataa
	def      A_				(			self							:   Optional[int]							):
						# there is always at least one resnet
						SCREAMING_SNAKE_CASE__													= [
						    FlaxResnetBlockaD(
						        in_channels=self.in_channels       ,   out_channels=self.in_channels       ,   dropout_prob=self.dropout       ,   dtype=self.dtype       ,   )
						]

						SCREAMING_SNAKE_CASE__													= []

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

											SCREAMING_SNAKE_CASE__													= FlaxResnetBlockaD(
											    in_channels=self.in_channels       ,   out_channels=self.in_channels       ,   dropout_prob=self.dropout       ,   dtype=self.dtype       ,   )
											resnets.append(UpperCAmelCase_							)

						SCREAMING_SNAKE_CASE__													= resnets
						SCREAMING_SNAKE_CASE__													= attentions





	def __call__(			self							:   str       ,   UpperCAmelCase_							:   Union[str, Any]       ,   UpperCAmelCase_							:   Optional[int]       ,   UpperCAmelCase_							:   List[str]       ,   UpperCAmelCase_							:   Any=True							):
						SCREAMING_SNAKE_CASE__													= self.resnets[0](UpperCAmelCase_       ,   UpperCAmelCase_							)
						for attn, resnet in zip(self.attentions       ,   self.resnets[1:]							):
											SCREAMING_SNAKE_CASE__													= attn(UpperCAmelCase_       ,   UpperCAmelCase_       ,   deterministic=UpperCAmelCase_							)
											SCREAMING_SNAKE_CASE__													= resnet(UpperCAmelCase_       ,   UpperCAmelCase_       ,   deterministic=UpperCAmelCase_							)

						return hidden_states





