
"""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,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)




@flax.struct.dataclass
class  __lowerCamelCase  ( lowerCAmelCase			):
				a__:      jnp.ndarray




@flax_register_to_config
class  __lowerCamelCase  ( nn.Module						,	lowerCAmelCase						,	lowerCAmelCase			):
				a__:      int          =    32
				a__:      int          =    4
				a__:      int          =    4
				a__:      Tuple[str]          =    (
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "DownBlock2D",
				)
				a__:      Tuple[str]          =    ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
				a__:      Union[bool, Tuple[bool]]          =    False
				a__:      Tuple[int]          =    (320, 640, 1280, 1280)
				a__:      int          =    2
				a__:      Union[int, Tuple[int]]          =    8
				a__:      Optional[Union[int, Tuple[int]]]          =    None
				a__:      int          =    1280
				a__:      float          =    0.0
				a__:      bool          =    False
				a__:      jnp.dtype          =    jnp.floataa
				a__:      bool          =    True
				a__:      int          =    0
				a__:      bool          =    False




				def        UpperCAmelCase__						(				self							,						UpperCAmelCase   ):
							# init input tensors
							lowerCamelCase_					=  (1, self.in_channels, self.sample_size, self.sample_size)
							lowerCamelCase_					=  jnp.zeros(UpperCAmelCase							,						dtype=jnp.floataa   )
							lowerCamelCase_					=  jnp.ones((1,)							,						dtype=jnp.intaa   )
							lowerCamelCase_					=  jnp.zeros((1, 1, self.cross_attention_dim)							,						dtype=jnp.floataa   )

							lowerCamelCase_		, lowerCamelCase_					=  jax.random.split(UpperCAmelCase   )
							lowerCamelCase_					=  {'''params''': params_rng, '''dropout''': dropout_rng}

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




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  self.block_out_channels
							lowerCamelCase_					=  block_out_channels[0] * 4

							if self.num_attention_heads is not None:
										raise ValueError(
										    '''At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.'''   )

							# 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.
							lowerCamelCase_					=  self.num_attention_heads or self.attention_head_dim

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

							# time
							lowerCamelCase_					=  FlaxTimesteps(
							    block_out_channels[0]							,						flip_sin_to_cos=self.flip_sin_to_cos							,						freq_shift=self.config.freq_shift   )
							lowerCamelCase_					=  FlaxTimestepEmbedding(UpperCAmelCase							,						dtype=self.dtype   )

							lowerCamelCase_					=  self.only_cross_attention
							if isinstance(UpperCAmelCase							,						UpperCAmelCase   ):
										lowerCamelCase_					=  (only_cross_attention,) * len(self.down_block_types   )

							if isinstance(UpperCAmelCase							,						UpperCAmelCase   ):
										lowerCamelCase_					=  (num_attention_heads,) * len(self.down_block_types   )

							# down
							lowerCamelCase_					=  []
							lowerCamelCase_					=  block_out_channels[0]
							for i, down_block_type in enumerate(self.down_block_types   ):
										lowerCamelCase_					=  output_channel
										lowerCamelCase_					=  block_out_channels[i]
										lowerCamelCase_					=  i == len(UpperCAmelCase   ) - 1

										if down_block_type == "CrossAttnDownBlock2D":
													lowerCamelCase_					=  FlaxCrossAttnDownBlockaD(
													    in_channels=UpperCAmelCase							,						out_channels=UpperCAmelCase							,						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]							,						use_memory_efficient_attention=self.use_memory_efficient_attention							,						dtype=self.dtype							,						)
										else:
													lowerCamelCase_					=  FlaxDownBlockaD(
													    in_channels=UpperCAmelCase							,						out_channels=UpperCAmelCase							,						dropout=self.dropout							,						num_layers=self.layers_per_block							,						add_downsample=not is_final_block							,						dtype=self.dtype							,						)

										down_blocks.append(UpperCAmelCase   )
							lowerCamelCase_					=  down_blocks

							# mid
							lowerCamelCase_					=  FlaxUNetMidBlockaDCrossAttn(
							    in_channels=block_out_channels[-1]							,						dropout=self.dropout							,						num_attention_heads=num_attention_heads[-1]							,						use_linear_projection=self.use_linear_projection							,						use_memory_efficient_attention=self.use_memory_efficient_attention							,						dtype=self.dtype							,						)

							# up
							lowerCamelCase_					=  []
							lowerCamelCase_					=  list(reversed(UpperCAmelCase   )   )
							lowerCamelCase_					=  list(reversed(UpperCAmelCase   )   )
							lowerCamelCase_					=  list(reversed(UpperCAmelCase   )   )
							lowerCamelCase_					=  reversed_block_out_channels[0]
							for i, up_block_type in enumerate(self.up_block_types   ):
										lowerCamelCase_					=  output_channel
										lowerCamelCase_					=  reversed_block_out_channels[i]
										lowerCamelCase_					=  reversed_block_out_channels[min(i + 1							,						len(UpperCAmelCase   ) - 1   )]

										lowerCamelCase_					=  i == len(UpperCAmelCase   ) - 1

										if up_block_type == "CrossAttnUpBlock2D":
													lowerCamelCase_					=  FlaxCrossAttnUpBlockaD(
													    in_channels=UpperCAmelCase							,						out_channels=UpperCAmelCase							,						prev_output_channel=UpperCAmelCase							,						num_layers=self.layers_per_block + 1							,						num_attention_heads=reversed_num_attention_heads[i]							,						add_upsample=not is_final_block							,						dropout=self.dropout							,						use_linear_projection=self.use_linear_projection							,						only_cross_attention=only_cross_attention[i]							,						use_memory_efficient_attention=self.use_memory_efficient_attention							,						dtype=self.dtype							,						)
										else:
													lowerCamelCase_					=  FlaxUpBlockaD(
													    in_channels=UpperCAmelCase							,						out_channels=UpperCAmelCase							,						prev_output_channel=UpperCAmelCase							,						num_layers=self.layers_per_block + 1							,						add_upsample=not is_final_block							,						dropout=self.dropout							,						dtype=self.dtype							,						)

										up_blocks.append(UpperCAmelCase   )
										lowerCamelCase_					=  output_channel
							lowerCamelCase_					=  up_blocks

							# out
							lowerCamelCase_					=  nn.GroupNorm(num_groups=32							,						epsilon=1e-5   )
							lowerCamelCase_					=  nn.Conv(
							    self.out_channels							,						kernel_size=(3, 3)							,						strides=(1, 1)							,						padding=((1, 1), (1, 1))							,						dtype=self.dtype							,						)





				def __call__(				self							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase=None							,						UpperCAmelCase=None							,						UpperCAmelCase = True							,						UpperCAmelCase = False							,						):
							# 1. time
							if not isinstance(UpperCAmelCase							,						jnp.ndarray   ):
										lowerCamelCase_					=  jnp.array([timesteps]							,						dtype=jnp.intaa   )
							elif isinstance(UpperCAmelCase							,						jnp.ndarray   ) and len(timesteps.shape   ) == 0:
										lowerCamelCase_					=  timesteps.astype(dtype=jnp.floataa   )
										lowerCamelCase_					=  jnp.expand_dims(UpperCAmelCase							,						0   )

							lowerCamelCase_					=  self.time_proj(UpperCAmelCase   )
							lowerCamelCase_					=  self.time_embedding(UpperCAmelCase   )

							# 2. pre-process
							lowerCamelCase_					=  jnp.transpose(UpperCAmelCase							,						(0, 2, 3, 1)   )
							lowerCamelCase_					=  self.conv_in(UpperCAmelCase   )

							# 3. down
							lowerCamelCase_					=  (sample,)
							for down_block in self.down_blocks:
										if isinstance(UpperCAmelCase							,						UpperCAmelCase   ):
													lowerCamelCase_		, lowerCamelCase_					=  down_block(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						deterministic=not train   )
										else:
													lowerCamelCase_		, lowerCamelCase_					=  down_block(UpperCAmelCase							,						UpperCAmelCase							,						deterministic=not train   )
										down_block_res_samples += res_samples

							if down_block_additional_residuals is not None:
										lowerCamelCase_					=  ()

										for down_block_res_sample, down_block_additional_residual in zip(
										    UpperCAmelCase							,						UpperCAmelCase   ):
													down_block_res_sample += down_block_additional_residual
													new_down_block_res_samples += (down_block_res_sample,)

										lowerCamelCase_					=  new_down_block_res_samples

							# 4. mid
							lowerCamelCase_					=  self.mid_block(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase							,						deterministic=not train   )

							if mid_block_additional_residual is not None:
										sample += mid_block_additional_residual

							# 5. up
							for up_block in self.up_blocks:
										lowerCamelCase_					=  down_block_res_samples[-(self.layers_per_block + 1) :]
										lowerCamelCase_					=  down_block_res_samples[: -(self.layers_per_block + 1)]
										if isinstance(UpperCAmelCase							,						UpperCAmelCase   ):
													lowerCamelCase_					=  up_block(
													    UpperCAmelCase							,						temb=UpperCAmelCase							,						encoder_hidden_states=UpperCAmelCase							,						res_hidden_states_tuple=UpperCAmelCase							,						deterministic=not train							,						)
										else:
													lowerCamelCase_					=  up_block(UpperCAmelCase							,						temb=UpperCAmelCase							,						res_hidden_states_tuple=UpperCAmelCase							,						deterministic=not train   )

        # 6. post-process
							lowerCamelCase_					=  self.conv_norm_out(UpperCAmelCase   )
							lowerCamelCase_					=  nn.silu(UpperCAmelCase   )
							lowerCamelCase_					=  self.conv_out(UpperCAmelCase   )
							lowerCamelCase_					=  jnp.transpose(UpperCAmelCase							,						(0, 3, 1, 2)   )

							if not return_dict:
										return (sample,)

							return FlaxUNetaDConditionOutput(sample=UpperCAmelCase   )





"""simple docstring"""


import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
						import optax
						from flax.training.common_utils import onehot

						from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
						from transformers.models.ta.modeling_flax_ta import shift_tokens_right




@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class  __lowerCamelCase  ( unittest.TestCase			):




				@slow
				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  FlaxMTaForConditionalGeneration.from_pretrained('''google/mt5-small'''   )
							lowerCamelCase_					=  AutoTokenizer.from_pretrained('''google/mt5-small'''   )

							lowerCamelCase_					=  tokenizer('''Hello there'''							,						return_tensors='''np'''   ).input_ids
							lowerCamelCase_					=  tokenizer('''Hi I am'''							,						return_tensors='''np'''   ).input_ids

							lowerCamelCase_					=  shift_tokens_right(UpperCAmelCase							,						model.config.pad_token_id							,						model.config.decoder_start_token_id   )

							lowerCamelCase_					=  model(UpperCAmelCase							,						decoder_input_ids=UpperCAmelCase   ).logits
							lowerCamelCase_					=  optax.softmax_cross_entropy(UpperCAmelCase							,						onehot(UpperCAmelCase							,						logits.shape[-1]   )   ).mean()

							lowerCamelCase_					=  -(labels.shape[-1] * loss.item())

							lowerCamelCase_					=  -8_4.9_1_2_7
							self.assertTrue(abs(mtf_score - EXPECTED_SCORE   ) < 1e-4   )



