
from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block

@dataclass
class    _a				(__lowerCamelCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE      =  42






class    _a				(nn.Module):





       """simple docstring"""




       def __init__(							self       ,							A__=3       ,							A__=3       ,							A__=("DownEncoderBlock2D",)       ,							A__=(64,)       ,							A__=2       ,							A__=32       ,							A__="silu"       ,							A__=True       ,							)    ->			Any:
          super().__init__()
          _SCREAMING_SNAKE_CASE      =				layers_per_block

          _SCREAMING_SNAKE_CASE      =				torch.nn.Convad(
              UpperCAmelCase_       ,							block_out_channels[0]       ,							kernel_size=3       ,							stride=1       ,							padding=1       ,							)

          _SCREAMING_SNAKE_CASE      =				None
          _SCREAMING_SNAKE_CASE      =				nn.ModuleList([]      )

          # down
          _SCREAMING_SNAKE_CASE      =				block_out_channels[0]
          for i, down_block_type in enumerate(UpperCAmelCase_      ):
             _SCREAMING_SNAKE_CASE      =				output_channel
             _SCREAMING_SNAKE_CASE      =				block_out_channels[i]
             _SCREAMING_SNAKE_CASE      =				i == len(UpperCAmelCase_      ) - 1

             _SCREAMING_SNAKE_CASE      =				get_down_block(
                 UpperCAmelCase_       ,							num_layers=self.layers_per_block       ,							in_channels=UpperCAmelCase_       ,							out_channels=UpperCAmelCase_       ,							add_downsample=not is_final_block       ,							resnet_eps=1E-6       ,							downsample_padding=0       ,							resnet_act_fn=UpperCAmelCase_       ,							resnet_groups=UpperCAmelCase_       ,							attention_head_dim=UpperCAmelCase_       ,							temb_channels=UpperCAmelCase_       ,							)
             self.down_blocks.append(UpperCAmelCase_      )

          # mid
          _SCREAMING_SNAKE_CASE      =				UNetMidBlockaD(
              in_channels=block_out_channels[-1]       ,							resnet_eps=1E-6       ,							resnet_act_fn=UpperCAmelCase_       ,							output_scale_factor=1       ,							resnet_time_scale_shift="""default"""       ,							attention_head_dim=block_out_channels[-1]       ,							resnet_groups=UpperCAmelCase_       ,							temb_channels=UpperCAmelCase_       ,							)

          # out
          _SCREAMING_SNAKE_CASE      =				nn.GroupNorm(num_channels=block_out_channels[-1]       ,							num_groups=UpperCAmelCase_       ,							eps=1E-6      )
          _SCREAMING_SNAKE_CASE      =				nn.SiLU()

          _SCREAMING_SNAKE_CASE      =				2 * out_channels if double_z else out_channels
          _SCREAMING_SNAKE_CASE      =				nn.Convad(block_out_channels[-1]       ,							UpperCAmelCase_       ,							3       ,							padding=1      )

          _SCREAMING_SNAKE_CASE      =				False






       def        UpperCamelCase				(							self       ,							A__      )    ->			Any:
          _SCREAMING_SNAKE_CASE      =				x
          _SCREAMING_SNAKE_CASE      =				self.conv_in(UpperCAmelCase_      )

          if self.training and self.gradient_checkpointing:

             def create_custom_forward(A__      ):
                def custom_forward(*A__      ):
                   return module(*UpperCAmelCase_      )

                return custom_forward

             # down
             if is_torch_version(""">="""       ,							"""1.11.0"""      ):
                for down_block in self.down_blocks:
                   _SCREAMING_SNAKE_CASE      =				torch.utils.checkpoint.checkpoint(
                       create_custom_forward(UpperCAmelCase_      )       ,							UpperCAmelCase_       ,							use_reentrant=UpperCAmelCase_      )
                # middle
                _SCREAMING_SNAKE_CASE      =				torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block      )       ,							UpperCAmelCase_       ,							use_reentrant=UpperCAmelCase_      )
             else:
                for down_block in self.down_blocks:
                   _SCREAMING_SNAKE_CASE      =				torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCAmelCase_      )       ,							UpperCAmelCase_      )
                # middle
                _SCREAMING_SNAKE_CASE      =				torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block      )       ,							UpperCAmelCase_      )

          else:
             # down
             for down_block in self.down_blocks:
                _SCREAMING_SNAKE_CASE      =				down_block(UpperCAmelCase_      )

             # middle
             _SCREAMING_SNAKE_CASE      =				self.mid_block(UpperCAmelCase_      )

          # post-process
          _SCREAMING_SNAKE_CASE      =				self.conv_norm_out(UpperCAmelCase_      )
          _SCREAMING_SNAKE_CASE      =				self.conv_act(UpperCAmelCase_      )
          _SCREAMING_SNAKE_CASE      =				self.conv_out(UpperCAmelCase_      )

          return sample






class    _a				(nn.Module):





       """simple docstring"""




       def __init__(							self       ,							A__=3       ,							A__=3       ,							A__=("UpDecoderBlock2D",)       ,							A__=(64,)       ,							A__=2       ,							A__=32       ,							A__="silu"       ,							A__="group"       ,							)    ->			int:
          super().__init__()
          _SCREAMING_SNAKE_CASE      =				layers_per_block

          _SCREAMING_SNAKE_CASE      =				nn.Convad(
              UpperCAmelCase_       ,							block_out_channels[-1]       ,							kernel_size=3       ,							stride=1       ,							padding=1       ,							)

          _SCREAMING_SNAKE_CASE      =				None
          _SCREAMING_SNAKE_CASE      =				nn.ModuleList([]      )

          _SCREAMING_SNAKE_CASE      =				in_channels if norm_type == 'spatial' else None

          # mid
          _SCREAMING_SNAKE_CASE      =				UNetMidBlockaD(
              in_channels=block_out_channels[-1]       ,							resnet_eps=1E-6       ,							resnet_act_fn=UpperCAmelCase_       ,							output_scale_factor=1       ,							resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type       ,							attention_head_dim=block_out_channels[-1]       ,							resnet_groups=UpperCAmelCase_       ,							temb_channels=UpperCAmelCase_       ,							)

          # up
          _SCREAMING_SNAKE_CASE      =				list(reversed(UpperCAmelCase_      )      )
          _SCREAMING_SNAKE_CASE      =				reversed_block_out_channels[0]
          for i, up_block_type in enumerate(UpperCAmelCase_      ):
             _SCREAMING_SNAKE_CASE      =				output_channel
             _SCREAMING_SNAKE_CASE      =				reversed_block_out_channels[i]

             _SCREAMING_SNAKE_CASE      =				i == len(UpperCAmelCase_      ) - 1

             _SCREAMING_SNAKE_CASE      =				get_up_block(
                 UpperCAmelCase_       ,							num_layers=self.layers_per_block + 1       ,							in_channels=UpperCAmelCase_       ,							out_channels=UpperCAmelCase_       ,							prev_output_channel=UpperCAmelCase_       ,							add_upsample=not is_final_block       ,							resnet_eps=1E-6       ,							resnet_act_fn=UpperCAmelCase_       ,							resnet_groups=UpperCAmelCase_       ,							attention_head_dim=UpperCAmelCase_       ,							temb_channels=UpperCAmelCase_       ,							resnet_time_scale_shift=UpperCAmelCase_       ,							)
             self.up_blocks.append(UpperCAmelCase_      )
             _SCREAMING_SNAKE_CASE      =				output_channel

          # out
          if norm_type == "spatial":
             _SCREAMING_SNAKE_CASE      =				SpatialNorm(block_out_channels[0]       ,							UpperCAmelCase_      )
          else:
             _SCREAMING_SNAKE_CASE      =				nn.GroupNorm(num_channels=block_out_channels[0]       ,							num_groups=UpperCAmelCase_       ,							eps=1E-6      )
          _SCREAMING_SNAKE_CASE      =				nn.SiLU()
          _SCREAMING_SNAKE_CASE      =				nn.Convad(block_out_channels[0]       ,							UpperCAmelCase_       ,							3       ,							padding=1      )

          _SCREAMING_SNAKE_CASE      =				False






       def        UpperCamelCase				(							self       ,							A__       ,							A__=None      )    ->			Dict:
          _SCREAMING_SNAKE_CASE      =				z
          _SCREAMING_SNAKE_CASE      =				self.conv_in(UpperCAmelCase_      )

          _SCREAMING_SNAKE_CASE      =				next(iter(self.up_blocks.parameters()      )      ).dtype
          if self.training and self.gradient_checkpointing:

             def create_custom_forward(A__      ):
                def custom_forward(*A__      ):
                   return module(*UpperCAmelCase_      )

                return custom_forward

             if is_torch_version(""">="""       ,							"""1.11.0"""      ):
                # middle
                _SCREAMING_SNAKE_CASE      =				torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block      )       ,							UpperCAmelCase_       ,							UpperCAmelCase_       ,							use_reentrant=UpperCAmelCase_      )
                _SCREAMING_SNAKE_CASE      =				sample.to(UpperCAmelCase_      )

                # up
                for up_block in self.up_blocks:
                   _SCREAMING_SNAKE_CASE      =				torch.utils.checkpoint.checkpoint(
                       create_custom_forward(UpperCAmelCase_      )       ,							UpperCAmelCase_       ,							UpperCAmelCase_       ,							use_reentrant=UpperCAmelCase_      )
             else:
                # middle
                _SCREAMING_SNAKE_CASE      =				torch.utils.checkpoint.checkpoint(
                    create_custom_forward(self.mid_block      )       ,							UpperCAmelCase_       ,							UpperCAmelCase_      )
                _SCREAMING_SNAKE_CASE      =				sample.to(UpperCAmelCase_      )

                # up
                for up_block in self.up_blocks:
                   _SCREAMING_SNAKE_CASE      =				torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCAmelCase_      )       ,							UpperCAmelCase_       ,							UpperCAmelCase_      )
          else:
             # middle
             _SCREAMING_SNAKE_CASE      =				self.mid_block(UpperCAmelCase_       ,							UpperCAmelCase_      )
             _SCREAMING_SNAKE_CASE      =				sample.to(UpperCAmelCase_      )

             # up
             for up_block in self.up_blocks:
                _SCREAMING_SNAKE_CASE      =				up_block(UpperCAmelCase_       ,							UpperCAmelCase_      )

        # post-process
          if latent_embeds is None:
             _SCREAMING_SNAKE_CASE      =				self.conv_norm_out(UpperCAmelCase_      )
          else:
             _SCREAMING_SNAKE_CASE      =				self.conv_norm_out(UpperCAmelCase_       ,							UpperCAmelCase_      )
          _SCREAMING_SNAKE_CASE      =				self.conv_act(UpperCAmelCase_      )
          _SCREAMING_SNAKE_CASE      =				self.conv_out(UpperCAmelCase_      )

          return sample






class    _a				(nn.Module):





       """simple docstring"""




       def __init__(							self       ,							A__       ,							A__       ,							A__       ,							A__=None       ,							A__="random"       ,							A__=False       ,							A__=True      )    ->			Optional[Any]:
          super().__init__()
          _SCREAMING_SNAKE_CASE      =				n_e
          _SCREAMING_SNAKE_CASE      =				vq_embed_dim
          _SCREAMING_SNAKE_CASE      =				beta
          _SCREAMING_SNAKE_CASE      =				legacy

          _SCREAMING_SNAKE_CASE      =				nn.Embedding(self.n_e       ,							self.vq_embed_dim      )
          self.embedding.weight.data.uniform_(-1.0 / self.n_e       ,							1.0 / self.n_e      )

          _SCREAMING_SNAKE_CASE      =				remap
          if self.remap is not None:
             self.register_buffer("""used"""       ,							torch.tensor(np.load(self.remap      )      )      )
             _SCREAMING_SNAKE_CASE      =				self.used.shape[0]
             _SCREAMING_SNAKE_CASE      =				unknown_index  # "random" or "extra" or integer
             if self.unknown_index == "extra":
                _SCREAMING_SNAKE_CASE      =				self.re_embed
                _SCREAMING_SNAKE_CASE      =				self.re_embed + 1
             print(
                 F"Remapping {self.n_e} indices to {self.re_embed} indices. "
                 F"Using {self.unknown_index} for unknown indices."      )
          else:
             _SCREAMING_SNAKE_CASE      =				n_e

          _SCREAMING_SNAKE_CASE      =				sane_index_shape






       def        UpperCamelCase				(							self       ,							A__      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				inds.shape
          assert len(UpperCAmelCase_      ) > 1
          _SCREAMING_SNAKE_CASE      =				inds.reshape(ishape[0]       ,							-1      )
          _SCREAMING_SNAKE_CASE      =				self.used.to(UpperCAmelCase_      )
          _SCREAMING_SNAKE_CASE      =				(inds[:, :, None] == used[None, None, ...]).long()
          _SCREAMING_SNAKE_CASE      =				match.argmax(-1      )
          _SCREAMING_SNAKE_CASE      =				match.sum(2      ) < 1
          if self.unknown_index == "random":
             _SCREAMING_SNAKE_CASE      =				torch.randint(0       ,							self.re_embed       ,							size=new[unknown].shape      ).to(device=new.device      )
          else:
             _SCREAMING_SNAKE_CASE      =				self.unknown_index
          return new.reshape(UpperCAmelCase_      )






       def        UpperCamelCase				(							self       ,							A__      )    ->			Optional[int]:
          _SCREAMING_SNAKE_CASE      =				inds.shape
          assert len(UpperCAmelCase_      ) > 1
          _SCREAMING_SNAKE_CASE      =				inds.reshape(ishape[0]       ,							-1      )
          _SCREAMING_SNAKE_CASE      =				self.used.to(UpperCAmelCase_      )
          if self.re_embed > self.used.shape[0]:  # extra token
             _SCREAMING_SNAKE_CASE      =				0  # simply set to zero
          _SCREAMING_SNAKE_CASE      =				torch.gather(used[None, :][inds.shape[0] * [0], :]       ,							1       ,							UpperCAmelCase_      )
          return back.reshape(UpperCAmelCase_      )






       def        UpperCamelCase				(							self       ,							A__      )    ->			List[Any]:
          # reshape z -> (batch, height, width, channel) and flatten
          _SCREAMING_SNAKE_CASE      =				z.permute(0       ,							2       ,							3       ,							1      ).contiguous()
          _SCREAMING_SNAKE_CASE      =				z.view(-1       ,							self.vq_embed_dim      )

          # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
          _SCREAMING_SNAKE_CASE      =				torch.argmin(torch.cdist(UpperCAmelCase_       ,							self.embedding.weight      )       ,							dim=1      )

          _SCREAMING_SNAKE_CASE      =				self.embedding(UpperCAmelCase_      ).view(z.shape      )
          _SCREAMING_SNAKE_CASE      =				None
          _SCREAMING_SNAKE_CASE      =				None

          # compute loss for embedding
          if not self.legacy:
             _SCREAMING_SNAKE_CASE      =				self.beta * torch.mean((z_q.detach() - z) ** 2      ) + torch.mean((z_q - z.detach()) ** 2      )
          else:
             _SCREAMING_SNAKE_CASE      =				torch.mean((z_q.detach() - z) ** 2      ) + self.beta * torch.mean((z_q - z.detach()) ** 2      )

          # preserve gradients
          _SCREAMING_SNAKE_CASE      =				z + (z_q - z).detach()

          # reshape back to match original input shape
          _SCREAMING_SNAKE_CASE      =				z_q.permute(0       ,							3       ,							1       ,							2      ).contiguous()

          if self.remap is not None:
             _SCREAMING_SNAKE_CASE      =				min_encoding_indices.reshape(z.shape[0]       ,							-1      )  # add batch axis
             _SCREAMING_SNAKE_CASE      =				self.remap_to_used(UpperCAmelCase_      )
             _SCREAMING_SNAKE_CASE      =				min_encoding_indices.reshape(-1       ,							1      )  # flatten

          if self.sane_index_shape:
             _SCREAMING_SNAKE_CASE      =				min_encoding_indices.reshape(z_q.shape[0]       ,							z_q.shape[2]       ,							z_q.shape[3]      )

          return z_q, loss, (perplexity, min_encodings, min_encoding_indices)






       def        UpperCamelCase				(							self       ,							A__       ,							A__      )    ->			Optional[Any]:
          # shape specifying (batch, height, width, channel)
          if self.remap is not None:
             _SCREAMING_SNAKE_CASE      =				indices.reshape(shape[0]       ,							-1      )  # add batch axis
             _SCREAMING_SNAKE_CASE      =				self.unmap_to_all(UpperCAmelCase_      )
             _SCREAMING_SNAKE_CASE      =				indices.reshape(-1      )  # flatten again

          # get quantized latent vectors
          _SCREAMING_SNAKE_CASE      =				self.embedding(UpperCAmelCase_      )

          if shape is not None:
             _SCREAMING_SNAKE_CASE      =				z_q.view(UpperCAmelCase_      )
             # reshape back to match original input shape
             _SCREAMING_SNAKE_CASE      =				z_q.permute(0       ,							3       ,							1       ,							2      ).contiguous()

          return z_q






class    _a				(__lowerCamelCase):





       """simple docstring"""




       def __init__(							self       ,							A__       ,							A__=False      )    ->			Optional[Any]:
          _SCREAMING_SNAKE_CASE      =				parameters
          _SCREAMING_SNAKE_CASE      =				torch.chunk(UpperCAmelCase_       ,							2       ,							dim=1      )
          _SCREAMING_SNAKE_CASE      =				torch.clamp(self.logvar       ,							-30.0       ,							20.0      )
          _SCREAMING_SNAKE_CASE      =				deterministic
          _SCREAMING_SNAKE_CASE      =				torch.exp(0.5 * self.logvar      )
          _SCREAMING_SNAKE_CASE      =				torch.exp(self.logvar      )
          if self.deterministic:
             _SCREAMING_SNAKE_CASE      =				torch.zeros_like(
                 self.mean       ,							device=self.parameters.device       ,							dtype=self.parameters.dtype      )






       def        UpperCamelCase				(							self       ,							A__ = None      )    ->			str:
          # make sure sample is on the same device as the parameters and has same dtype
          _SCREAMING_SNAKE_CASE      =				randn_tensor(
              self.mean.shape       ,							generator=UpperCAmelCase_       ,							device=self.parameters.device       ,							dtype=self.parameters.dtype      )
          _SCREAMING_SNAKE_CASE      =				self.mean + self.std * sample
          return x






       def        UpperCamelCase				(							self       ,							A__=None      )    ->			Optional[int]:
          if self.deterministic:
             return torch.Tensor([0.0]      )
          else:
             if other is None:
                return 0.5 * torch.sum(torch.pow(self.mean       ,							2      ) + self.var - 1.0 - self.logvar       ,							dim=[1, 2, 3]      )
             else:
                return 0.5 * torch.sum(
                    torch.pow(self.mean - other.mean       ,							2      ) / other.var
                    + self.var / other.var
                    - 1.0
                    - self.logvar
                    + other.logvar       ,							dim=[1, 2, 3]       ,							)






       def        UpperCamelCase				(							self       ,							A__       ,							A__=[1, 2, 3]      )    ->			Union[str, Any]:
          if self.deterministic:
             return torch.Tensor([0.0]      )
          _SCREAMING_SNAKE_CASE      =				np.log(2.0 * np.pi      )
          return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean       ,							2      ) / self.var       ,							dim=UpperCAmelCase_      )






       def        UpperCamelCase				(							self      )    ->			Tuple:
          return self.mean








'''simple docstring'''



from sklearn.metrics import matthews_corrcoef

import datasets


UpperCamelCase__      :		List[str]       							=  "\nCompute the Matthews correlation coefficient (MCC)\n\nThe Matthews correlation coefficient is used in machine learning as a\nmeasure of the quality of binary and multiclass classifications. It takes\ninto account true and false positives and negatives and is generally\nregarded as a balanced measure which can be used even if the classes are of\nvery different sizes. The MCC is in essence a correlation coefficient value\nbetween -1 and +1. A coefficient of +1 represents a perfect prediction, 0\nan average random prediction and -1 an inverse prediction.  The statistic\nis also known as the phi coefficient. [source: Wikipedia]\n"

UpperCamelCase__      :		List[Any]       							=  "\nArgs:\n    predictions (list of int): Predicted labels, as returned by a model.\n    references (list of int): Ground truth labels.\n    sample_weight (list of int, float, or bool): Sample weights. Defaults to `None`.\nReturns:\n    matthews_correlation (dict containing float): Matthews correlation.\nExamples:\n    Example 1, a basic example with only predictions and references as inputs:\n        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")\n        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],\n        ...                                     predictions=[1, 2, 2, 0, 3, 3])\n        >>> print(round(results['matthews_correlation'], 2))\n        0.54\n\n    Example 2, the same example as above, but also including sample weights:\n        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")\n        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],\n        ...                                     predictions=[1, 2, 2, 0, 3, 3],\n        ...                                     sample_weight=[0.5, 3, 1, 1, 1, 2])\n        >>> print(round(results['matthews_correlation'], 2))\n        0.1\n\n    Example 3, the same example as above, but with sample weights that cause a negative correlation:\n        >>> matthews_metric = datasets.load_metric(\"matthews_correlation\")\n        >>> results = matthews_metric.compute(references=[1, 3, 2, 0, 3, 2],\n        ...                                     predictions=[1, 2, 2, 0, 3, 3],\n        ...                                     sample_weight=[0.5, 1, 0, 0, 0, 1])\n        >>> print(round(results['matthews_correlation'], 2))\n        -0.25\n"

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

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





       """simple docstring"""




       def        UpperCamelCase				(							self      )    ->			Optional[Any]:
          return datasets.MetricInfo(
              description=_DESCRIPTION       ,							citation=_CITATION       ,							inputs_description=_KWARGS_DESCRIPTION       ,							features=datasets.Features(
                  {
                      """predictions""": datasets.Value("""int32"""      ),
                      """references""": datasets.Value("""int32"""      ),
                  }      )       ,							reference_urls=[
                  """https://scikit-learn.org/stable/modules/generated/sklearn.metrics.matthews_corrcoef.html"""
              ]       ,							)






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__=None      )    ->			List[str]:
          return {
              "matthews_correlation": float(matthews_corrcoef(A__       ,							A__       ,							sample_weight=A__      )      ),
          }






