




'''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    __lowercase						(  __magic_name__ ):
   _a							=      42





@flax_register_to_config
class    __lowercase						(  nn.Module       ,							__magic_name__       ,							__magic_name__ ):
   _a							=      32
   _a							=      4
   _a							=      4
   _a							=      (
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "DownBlock2D",
   )
   _a							=      ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
   _a							=      False
   _a							=      (320, 640, 1280, 1280)
   _a							=      2
   _a							=      8
   _a							=      None
   _a							=      1280
   _a							=      0.0
   _a							=      False
   _a							=      jnp.floataa
   _a							=      True
   _a							=      0
   _a							=      False







   def 		UpperCamelCase__			(	self      ,							UpperCamelCase      )				->    FrozenDict:
     # init input tensors
     __a			=					(1, self.in_channels, self.sample_size, self.sample_size)
     __a			=					jnp.zeros(UpperCamelCase      ,							dtype=jnp.floataa      )
     __a			=					jnp.ones((1,)      ,							dtype=jnp.intaa      )
     __a			=					jnp.zeros((1, 1, self.cross_attention_dim)      ,							dtype=jnp.floataa      )

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

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







   def 		UpperCamelCase__			(	self      )				->    int:
     __a			=					self.block_out_channels
     __a			=					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.
     __a			=					self.num_attention_heads or self.attention_head_dim

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

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

     __a			=					self.only_cross_attention
     if isinstance(UpperCamelCase      ,							UpperCamelCase      ):
       __a			=					(only_cross_attention,) * len(self.down_block_types      )

     if isinstance(UpperCamelCase      ,							UpperCamelCase      ):
       __a			=					(num_attention_heads,) * len(self.down_block_types      )

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

       if down_block_type == "CrossAttnDownBlock2D":
         __a			=					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:
         __a			=					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      )
     __a			=					down_blocks

     # mid
     __a			=					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
     __a			=					[]
     __a			=					list(reversed(UpperCamelCase      )      )
     __a			=					list(reversed(UpperCamelCase      )      )
     __a			=					list(reversed(UpperCamelCase      )      )
     __a			=					reversed_block_out_channels[0]
     for i, up_block_type in enumerate(self.up_block_types      ):
       __a			=					output_channel
       __a			=					reversed_block_out_channels[i]
       __a			=					reversed_block_out_channels[min(i + 1      ,							len(UpperCamelCase      ) - 1      )]

       __a			=					i == len(UpperCamelCase      ) - 1

       if up_block_type == "CrossAttnUpBlock2D":
         __a			=					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:
         __a			=					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      )
       __a			=					output_channel
     __a			=					up_blocks

     # out
     __a			=					nn.GroupNorm(num_groups=32      ,							epsilon=1e-5      )
     __a			=					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      ,							)				->    Union[FlaxUNetaDConditionOutput, Tuple]:
     # 1. time
     if not isinstance(UpperCamelCase      ,							jnp.ndarray      ):
       __a			=					jnp.array([timesteps]      ,							dtype=jnp.intaa      )
     elif isinstance(UpperCamelCase      ,							jnp.ndarray      ) and len(timesteps.shape      ) == 0:
       __a			=					timesteps.astype(dtype=jnp.floataa      )
       __a			=					jnp.expand_dims(UpperCamelCase      ,							0      )

     __a			=					self.time_proj(UpperCamelCase      )
     __a			=					self.time_embedding(UpperCamelCase      )

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

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

     if down_block_additional_residuals is not None:
       __a			=					()

       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,)

       __a			=					new_down_block_res_samples

     # 4. mid
     __a			=					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:
       __a			=					down_block_res_samples[-(self.layers_per_block + 1) :]
       __a			=					down_block_res_samples[: -(self.layers_per_block + 1)]
       if isinstance(UpperCamelCase      ,							UpperCamelCase      ):
         __a			=					up_block(
             UpperCamelCase      ,							temb=UpperCamelCase      ,							encoder_hidden_states=UpperCamelCase      ,							res_hidden_states_tuple=UpperCamelCase      ,							deterministic=not train      ,							)
       else:
         __a			=					up_block(UpperCamelCase      ,							temb=UpperCamelCase      ,							res_hidden_states_tuple=UpperCamelCase      ,							deterministic=not train      )

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

     if not return_dict:
       return (sample,)

     return FlaxUNetaDConditionOutput(sample=UpperCamelCase      )









'''simple docstring'''


import numpy as np

import datasets


UpperCAmelCase_ 		=      "\nCompute the Mahalanobis Distance\n\nMahalonobis distance is the distance between a point and a distribution.\nAnd not between two distinct points. It is effectively a multivariate equivalent of the Euclidean distance.\nIt was introduced by Prof. P. C. Mahalanobis in 1936\nand has been used in various statistical applications ever since\n[source: https://www.machinelearningplus.com/statistics/mahalanobis-distance/]\n"

UpperCAmelCase_ 		=      "\\n@article{de2000mahalanobis,\n  title={The mahalanobis distance},\n  author={De Maesschalck, Roy and Jouan-Rimbaud, Delphine and Massart, D{\'e}sir{\'e} L},\n  journal={Chemometrics and intelligent laboratory systems},\n  volume={50},\n  number={1},\n  pages={1--18},\n  year={2000},\n  publisher={Elsevier}\n}\n"

UpperCAmelCase_ 		=      "\nArgs:\n    X: List of datapoints to be compared with the `reference_distribution`.\n    reference_distribution: List of datapoints from the reference distribution we want to compare to.\nReturns:\n    mahalanobis: The Mahalonobis distance for each datapoint in `X`.\nExamples:\n\n    >>> mahalanobis_metric = datasets.load_metric(\"mahalanobis\")\n    >>> results = mahalanobis_metric.compute(reference_distribution=[[0, 1], [1, 0]], X=[[0, 1]])\n    >>> print(results)\n    {'mahalanobis': array([0.5])}\n"





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







   def 		UpperCamelCase__			(	self      )				->    Optional[Any]:
     return datasets.MetricInfo(
         description=_DESCRIPTION      ,							citation=_CITATION      ,							inputs_description=_KWARGS_DESCRIPTION      ,							features=datasets.Features(
             {
                 'X': datasets.Sequence(datasets.Value('float'      ,							id='sequence'      )      ,							id='X'      ),
             }      )      ,							)







   def 		UpperCamelCase__			(	self      ,							UpperCamelCase      ,							UpperCamelCase      )				->    Optional[int]:
     # convert to numpy arrays
     __a			=					np.array(UpperCamelCase      )
     __a			=					np.array(UpperCamelCase      )

     # Assert that arrays are 2D
     if len(X.shape      ) != 2:
       raise ValueError('Expected `X` to be a 2D vector'      )
     if len(reference_distribution.shape      ) != 2:
       raise ValueError('Expected `reference_distribution` to be a 2D vector'      )
     if reference_distribution.shape[0] < 2:
       raise ValueError(
           'Expected `reference_distribution` to be a 2D vector with more than one element in the first dimension'      )

     # Get mahalanobis distance for each prediction
     __a			=					X - np.mean(UpperCamelCase      )
     __a			=					np.cov(reference_distribution.T      )
     try:
       __a			=					np.linalg.inv(UpperCamelCase      )
     except np.linalg.LinAlgError:
       __a			=					np.linalg.pinv(UpperCamelCase      )
     __a			=					np.dot(UpperCamelCase      ,							UpperCamelCase      )
     __a			=					np.dot(UpperCamelCase      ,							X_minus_mu.T      ).diagonal()

     return {"mahalanobis": mahal_dist}



