



'''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 		_UpperCamelCase						(							lowerCamelCase__					):


			'''simple docstring'''

			_A   :	jnp.ndarray






@flax_register_to_config
class 		_UpperCamelCase						(							nn.Module					,      lowerCamelCase__					,      lowerCamelCase__					):


			'''simple docstring'''

			_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  :				Dict							,  lowerCAmelCase__  :				jax.random.KeyArray		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	List[str]					 =	(1, self.in_channels, self.sample_size, self.sample_size)
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	jnp.zeros(lowerCAmelCase__							,  dtype=jnp.floataa		)
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	jnp.ones((1,)							,  dtype=jnp.intaa		)
					__SCREAMING_SNAKE_CASE :	int					 =	jnp.zeros((1, 1, self.cross_attention_dim)							,  dtype=jnp.floataa		)

					__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Tuple					 =	jax.random.split(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	{"""params""": params_rng, """dropout""": dropout_rng}

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

			def UpperCamelCase__      (							self  :				Dict		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	int					 =	self.block_out_channels
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	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.
					__SCREAMING_SNAKE_CASE :	Any					 =	self.num_attention_heads or self.attention_head_dim

					# input
					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	nn.Conv(
					    block_out_channels[0]							,  kernel_size=(3, 3)							,  strides=(1, 1)							,  padding=((1, 1), (1, 1))							,  dtype=self.dtype							,  )

					# time
					__SCREAMING_SNAKE_CASE :	Dict					 =	FlaxTimesteps(
					    block_out_channels[0]							,  flip_sin_to_cos=self.flip_sin_to_cos							,  freq_shift=self.config.freq_shift		)
					__SCREAMING_SNAKE_CASE :	int					 =	FlaxTimestepEmbedding(lowerCAmelCase__							,  dtype=self.dtype		)

					__SCREAMING_SNAKE_CASE :	Dict					 =	self.only_cross_attention
					if isinstance(lowerCAmelCase__							,  lowerCAmelCase__		):
							__SCREAMING_SNAKE_CASE :	int					 =	(only_cross_attention,) * len(self.down_block_types		)

					if isinstance(lowerCAmelCase__							,  lowerCAmelCase__		):
							__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	(num_attention_heads,) * len(self.down_block_types		)

					# down
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	[]
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	block_out_channels[0]
					for i, down_block_type in enumerate(self.down_block_types		):
							__SCREAMING_SNAKE_CASE :	Optional[int]					 =	output_channel
							__SCREAMING_SNAKE_CASE :	Optional[int]					 =	block_out_channels[i]
							__SCREAMING_SNAKE_CASE :	List[Any]					 =	i == len(lowerCAmelCase__		) - 1

							if down_block_type == "CrossAttnDownBlock2D":
									__SCREAMING_SNAKE_CASE :	List[Any]					 =	FlaxCrossAttnDownBlockaD(
									    in_channels=lowerCAmelCase__							,  out_channels=lowerCAmelCase__							,  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:
									__SCREAMING_SNAKE_CASE :	Optional[int]					 =	FlaxDownBlockaD(
									    in_channels=lowerCAmelCase__							,  out_channels=lowerCAmelCase__							,  dropout=self.dropout							,  num_layers=self.layers_per_block							,  add_downsample=not is_final_block							,  dtype=self.dtype							,  )

							down_blocks.append(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	down_blocks

					# mid
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	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
					__SCREAMING_SNAKE_CASE :	str					 =	[]
					__SCREAMING_SNAKE_CASE :	Tuple					 =	list(reversed(lowerCAmelCase__		)		)
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	list(reversed(lowerCAmelCase__		)		)
					__SCREAMING_SNAKE_CASE :	int					 =	list(reversed(lowerCAmelCase__		)		)
					__SCREAMING_SNAKE_CASE :	List[str]					 =	reversed_block_out_channels[0]
					for i, up_block_type in enumerate(self.up_block_types		):
							__SCREAMING_SNAKE_CASE :	Tuple					 =	output_channel
							__SCREAMING_SNAKE_CASE :	Any					 =	reversed_block_out_channels[i]
							__SCREAMING_SNAKE_CASE :	str					 =	reversed_block_out_channels[min(i + 1							,  len(lowerCAmelCase__		) - 1		)]

							__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	i == len(lowerCAmelCase__		) - 1

							if up_block_type == "CrossAttnUpBlock2D":
									__SCREAMING_SNAKE_CASE :	Optional[int]					 =	FlaxCrossAttnUpBlockaD(
									    in_channels=lowerCAmelCase__							,  out_channels=lowerCAmelCase__							,  prev_output_channel=lowerCAmelCase__							,  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:
									__SCREAMING_SNAKE_CASE :	Any					 =	FlaxUpBlockaD(
									    in_channels=lowerCAmelCase__							,  out_channels=lowerCAmelCase__							,  prev_output_channel=lowerCAmelCase__							,  num_layers=self.layers_per_block + 1							,  add_upsample=not is_final_block							,  dropout=self.dropout							,  dtype=self.dtype							,  )

							up_blocks.append(lowerCAmelCase__		)
							__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	output_channel
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	up_blocks

					# out
					__SCREAMING_SNAKE_CASE :	str					 =	nn.GroupNorm(num_groups=3_2							,  epsilon=1E-5		)
					__SCREAMING_SNAKE_CASE :	Dict					 =	nn.Conv(
					    self.out_channels							,  kernel_size=(3, 3)							,  strides=(1, 1)							,  padding=((1, 1), (1, 1))							,  dtype=self.dtype							,  )



			def __call__(							self  :				Tuple							,  lowerCAmelCase__  :				Union[str, Any]							,  lowerCAmelCase__  :				Dict							,  lowerCAmelCase__  :				Optional[int]							,  lowerCAmelCase__  :				Tuple=None							,  lowerCAmelCase__  :				List[str]=None							,  lowerCAmelCase__  :				bool = True							,  lowerCAmelCase__  :				bool = False							,  ):





					"""simple docstring"""






					if not isinstance(lowerCAmelCase__							,  jnp.ndarray		):
							__SCREAMING_SNAKE_CASE :	Dict					 =	jnp.array([timesteps]							,  dtype=jnp.intaa		)
					elif isinstance(lowerCAmelCase__							,  jnp.ndarray		) and len(timesteps.shape		) == 0:
							__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	timesteps.astype(dtype=jnp.floataa		)
							__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	jnp.expand_dims(lowerCAmelCase__							,  0		)

					__SCREAMING_SNAKE_CASE :	Any					 =	self.time_proj(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Dict					 =	self.time_embedding(lowerCAmelCase__		)

					# 2. pre-process
					__SCREAMING_SNAKE_CASE :	List[str]					 =	jnp.transpose(lowerCAmelCase__							,  (0, 2, 3, 1)		)
					__SCREAMING_SNAKE_CASE :	str					 =	self.conv_in(lowerCAmelCase__		)

					# 3. down
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	(sample,)
					for down_block in self.down_blocks:
							if isinstance(lowerCAmelCase__							,  lowerCAmelCase__		):
									__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	List[Any]					 =	down_block(lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  deterministic=not train		)
							else:
									__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	List[str]					 =	down_block(lowerCAmelCase__							,  lowerCAmelCase__							,  deterministic=not train		)
							down_block_res_samples += res_samples

					if down_block_additional_residuals is not None:
							__SCREAMING_SNAKE_CASE :	Tuple					 =	()

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

							__SCREAMING_SNAKE_CASE :	int					 =	new_down_block_res_samples

					# 4. mid
					__SCREAMING_SNAKE_CASE :	Dict					 =	self.mid_block(lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  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:
							__SCREAMING_SNAKE_CASE :	str					 =	down_block_res_samples[-(self.layers_per_block + 1) :]
							__SCREAMING_SNAKE_CASE :	int					 =	down_block_res_samples[: -(self.layers_per_block + 1)]
							if isinstance(lowerCAmelCase__							,  lowerCAmelCase__		):
									__SCREAMING_SNAKE_CASE :	Dict					 =	up_block(
									    lowerCAmelCase__							,  temb=lowerCAmelCase__							,  encoder_hidden_states=lowerCAmelCase__							,  res_hidden_states_tuple=lowerCAmelCase__							,  deterministic=not train							,  )
							else:
									__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	up_block(lowerCAmelCase__							,  temb=lowerCAmelCase__							,  res_hidden_states_tuple=lowerCAmelCase__							,  deterministic=not train		)

        # 6. post-process
					__SCREAMING_SNAKE_CASE :	Tuple					 =	self.conv_norm_out(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	str					 =	nn.silu(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Tuple					 =	self.conv_out(lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	int					 =	jnp.transpose(lowerCAmelCase__							,  (0, 3, 1, 2)		)

					if not return_dict:
							return (sample,)

					return FlaxUNetaDConditionOutput(sample=lowerCAmelCase__		)




'''simple docstring'''




from sklearn.metrics import mean_squared_error

import datasets


UpperCamelCase__ :			Optional[Any]												=				'''\
@article{scikit-learn,
 title={Scikit-learn: Machine Learning in {P}ython},
 author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
 journal={Journal of Machine Learning Research},
 volume={12},
 pages={2825--2830},
 year={2011}
}
'''

UpperCamelCase__ :			List[str]												=				'''\
Mean Squared Error(MSE) is the average of the square of difference between the predicted
and actual values.
'''


UpperCamelCase__ :			Dict												=				'''
Args:
    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Estimated target values.
    references: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Ground truth (correct) target values.
    sample_weight: array-like of shape (n_samples,), default=None
        Sample weights.
    multioutput: {"raw_values", "uniform_average"} or array-like of shape (n_outputs,), default="uniform_average"
        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.

        "raw_values" : Returns a full set of errors in case of multioutput input.

        "uniform_average" : Errors of all outputs are averaged with uniform weight.

    squared : bool, default=True
        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.

Returns:
    mse : mean squared error.
Examples:

    >>> mse_metric = datasets.load_metric("mse")
    >>> predictions = [2.5, 0.0, 2, 8]
    >>> references = [3, -0.5, 2, 7]
    >>> results = mse_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'mse\': 0.375}
    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)
    >>> print(rmse_result)
    {\'mse\': 0.6123724356957945}

    If you\'re using multi-dimensional lists, then set the config as follows :

    >>> mse_metric = datasets.load_metric("mse", "multilist")
    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]
    >>> references = [[0, 2], [-1, 2], [8, -5]]
    >>> results = mse_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'mse\': 0.7083333333333334}
    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput=\'raw_values\')
    >>> print(results) # doctest: +NORMALIZE_WHITESPACE
    {\'mse\': array([0.41666667, 1. ])}
'''






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION					,      _KWARGS_DESCRIPTION					)
class 		_UpperCamelCase						(							datasets.Metric					):


			'''simple docstring'''

			def UpperCamelCase__      (							self  :				str		):





					"""simple docstring"""






					return datasets.MetricInfo(
					    description=_DESCRIPTION							,  citation=_CITATION							,  inputs_description=_KWARGS_DESCRIPTION							,  features=datasets.Features(self._get_feature_types()		)							,  reference_urls=[
					        """https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html"""
					    ]							,  )

			def UpperCamelCase__      (							self  :				Optional[Any]		):





					"""simple docstring"""






					if self.config_name == "multilist":
							return {
							    "predictions": datasets.Sequence(datasets.Value("""float"""		)		),
							    "references": datasets.Sequence(datasets.Value("""float"""		)		),
							}
					else:
							return {
							    "predictions": datasets.Value("""float"""		),
							    "references": datasets.Value("""float"""		),
							}



			def UpperCamelCase__      (							self  :				List[Any]							,  lowerCAmelCase__  :				Any							,  lowerCAmelCase__  :				str							,  lowerCAmelCase__  :				List[Any]=None							,  lowerCAmelCase__  :				Optional[int]="uniform_average"							,  lowerCAmelCase__  :				int=True		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	str					 =	mean_squared_error(
					    lowerCAmelCase__							,  lowerCAmelCase__							,  sample_weight=lowerCAmelCase__							,  multioutput=lowerCAmelCase__							,  squared=lowerCAmelCase__		)

					return {"mse": mse}