
from collections import defaultdict

from graphs.minimum_spanning_tree_prims import prisms_algorithm as mst

def a__    (       ):






  '''simple docstring'''





  __magic_name__     ,    __magic_name__				     =       9, 14  # noqa: F841
  __magic_name__				     =       [
      [0, 1, 4],
      [0, 7, 8],
      [1, 2, 8],
      [7, 8, 7],
      [7, 6, 1],
      [2, 8, 2],
      [8, 6, 6],
      [2, 3, 7],
      [2, 5, 4],
      [6, 5, 2],
      [3, 5, 14],
      [3, 4, 9],
      [5, 4, 10],
      [1, 7, 11],
  ]

  __magic_name__				     =       defaultdict(A_					)
  for nodea, nodea, cost in edges:
    adjancency[nodea].append([nodea, cost]					)
    adjancency[nodea].append([nodea, cost]					)

  __magic_name__				     =       mst(A_					)

  __magic_name__				     =       [
      [7, 6, 1],
      [2, 8, 2],
      [6, 5, 2],
      [0, 1, 4],
      [2, 5, 4],
      [2, 3, 7],
      [0, 7, 8],
      [3, 4, 9],
  ]

  for answer in expected:
    __magic_name__				     =       tuple(answer[:2]					)
    __magic_name__				     =       tuple(edge[::-1]					)
    assert edge in result or reverse in result



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        UpperCAmelCase_			( _A     ):


  '''simple docstring'''







  a__           =		42







class        UpperCAmelCase_			( nn.Module     ):


  '''simple docstring'''







  def __init__( self			: Tuple			,							UpperCamelCase__			: List[Any]=3			,							UpperCamelCase__			: str=3			,							UpperCamelCase__			: List[Any]=("DownEncoderBlock2D",)			,							UpperCamelCase__			: Optional[Any]=(64,)			,							UpperCamelCase__			: Optional[Any]=2			,							UpperCamelCase__			: Union[str, Any]=32			,							UpperCamelCase__			: Optional[Any]="silu"			,							UpperCamelCase__			: List[str]=True			,							)			->	str:







    """simple docstring"""






    super().__init__()
    __magic_name__				     =       layers_per_block

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

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

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

      __magic_name__				     =       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
    __magic_name__				     =       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
    __magic_name__				     =       nn.GroupNorm(num_channels=block_out_channels[-1]			,							num_groups=UpperCamelCase__			,							eps=1E-6    )
    __magic_name__				     =       nn.SiLU()

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

    __magic_name__				     =       False





  def 							_lowercase							( self			: List[str]			,							UpperCamelCase__			: Optional[Any]    )			->	int:







    """simple docstring"""






    __magic_name__				     =       x
    __magic_name__				     =       self.conv_in(UpperCamelCase__    )

    if self.training and self.gradient_checkpointing:

      def create_custom_forward(UpperCamelCase__			: int    ):
        def custom_forward(*UpperCamelCase__			: str    ):
          return module(*UpperCamelCase__    )

        return custom_forward

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

    else:
      # down
      for down_block in self.down_blocks:
        __magic_name__				     =       down_block(UpperCamelCase__    )

      # middle
      __magic_name__				     =       self.mid_block(UpperCamelCase__    )

    # post-process
    __magic_name__				     =       self.conv_norm_out(UpperCamelCase__    )
    __magic_name__				     =       self.conv_act(UpperCamelCase__    )
    __magic_name__				     =       self.conv_out(UpperCamelCase__    )

    return sample







class        UpperCAmelCase_			( nn.Module     ):


  '''simple docstring'''







  def __init__( self			: List[str]			,							UpperCamelCase__			: int=3			,							UpperCamelCase__			: Dict=3			,							UpperCamelCase__			: List[Any]=("UpDecoderBlock2D",)			,							UpperCamelCase__			: List[Any]=(64,)			,							UpperCamelCase__			: Dict=2			,							UpperCamelCase__			: int=32			,							UpperCamelCase__			: Optional[int]="silu"			,							UpperCamelCase__			: Tuple="group"			,							)			->	Dict:







    """simple docstring"""






    super().__init__()
    __magic_name__				     =       layers_per_block

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

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

    __magic_name__				     =       in_channels if norm_type == """spatial""" else None

    # mid
    __magic_name__				     =       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
    __magic_name__				     =       list(reversed(UpperCamelCase__    )    )
    __magic_name__				     =       reversed_block_out_channels[0]
    for i, up_block_type in enumerate(UpperCamelCase__    ):
      __magic_name__				     =       output_channel
      __magic_name__				     =       reversed_block_out_channels[i]

      __magic_name__				     =       i == len(UpperCamelCase__    ) - 1

      __magic_name__				     =       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__    )
      __magic_name__				     =       output_channel

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

    __magic_name__				     =       False





  def 							_lowercase							( self			: str			,							UpperCamelCase__			: Union[str, Any]			,							UpperCamelCase__			: Tuple=None    )			->	Tuple:







    """simple docstring"""






    __magic_name__				     =       z
    __magic_name__				     =       self.conv_in(UpperCamelCase__    )

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

      def create_custom_forward(UpperCamelCase__			: Optional[int]    ):
        def custom_forward(*UpperCamelCase__			: int    ):
          return module(*UpperCamelCase__    )

        return custom_forward

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

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

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

      # up
      for up_block in self.up_blocks:
        __magic_name__				     =       up_block(UpperCamelCase__			,							UpperCamelCase__    )

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

    return sample







class        UpperCAmelCase_			( nn.Module     ):


  '''simple docstring'''







  def __init__( self			: str			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: List[str]			,							UpperCamelCase__			: Union[str, Any]=None			,							UpperCamelCase__			: Dict="random"			,							UpperCamelCase__			: List[Any]=False			,							UpperCamelCase__			: Dict=True    )			->	Optional[Any]:







    """simple docstring"""






    super().__init__()
    __magic_name__				     =       n_e
    __magic_name__				     =       vq_embed_dim
    __magic_name__				     =       beta
    __magic_name__				     =       legacy

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

    __magic_name__				     =       remap
    if self.remap is not None:
      self.register_buffer("""used"""			,							torch.tensor(np.load(self.remap    )    )    )
      __magic_name__				     =       self.used.shape[0]
      __magic_name__				     =       unknown_index  # "random" or "extra" or integer
      if self.unknown_index == "extra":
        __magic_name__				     =       self.re_embed
        __magic_name__				     =       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:
      __magic_name__				     =       n_e

    __magic_name__				     =       sane_index_shape
  def 							_lowercase							( self			: Union[str, Any]			,							UpperCamelCase__			: Dict    )			->	Union[str, Any]:







    """simple docstring"""






    __magic_name__				     =       inds.shape
    assert len(UpperCamelCase__    ) > 1
    __magic_name__				     =       inds.reshape(ishape[0]			,							-1    )
    __magic_name__				     =       self.used.to(UpperCamelCase__    )
    __magic_name__				     =       (inds[:, :, None] == used[None, None, ...]).long()
    __magic_name__				     =       match.argmax(-1    )
    __magic_name__				     =       match.sum(2    ) < 1
    if self.unknown_index == "random":
      __magic_name__				     =       torch.randint(0			,							self.re_embed			,							size=new[unknown].shape    ).to(device=new.device    )
    else:
      __magic_name__				     =       self.unknown_index
    return new.reshape(UpperCamelCase__    )
  def 							_lowercase							( self			: Optional[Any]			,							UpperCamelCase__			: str    )			->	Tuple:







    """simple docstring"""






    __magic_name__				     =       inds.shape
    assert len(UpperCamelCase__    ) > 1
    __magic_name__				     =       inds.reshape(ishape[0]			,							-1    )
    __magic_name__				     =       self.used.to(UpperCamelCase__    )
    if self.re_embed > self.used.shape[0]:  # extra token
      __magic_name__				     =       0  # simply set to zero
    __magic_name__				     =       torch.gather(used[None, :][inds.shape[0] * [0], :]			,							1			,							UpperCamelCase__    )
    return back.reshape(UpperCamelCase__    )
  def 							_lowercase							( self			: List[str]			,							UpperCamelCase__			: List[str]    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       z.permute(0			,							2			,							3			,							1    ).contiguous()
    __magic_name__				     =       z.view(-1			,							self.vq_embed_dim    )

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

    __magic_name__				     =       self.embedding(UpperCamelCase__    ).view(z.shape    )
    __magic_name__				     =       None
    __magic_name__				     =       None

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

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

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

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

    if self.sane_index_shape:
      __magic_name__				     =       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 							_lowercase							( self			: Union[str, Any]			,							UpperCamelCase__			: List[str]			,							UpperCamelCase__			: Union[str, Any]    )			->	int:







    """simple docstring"""






    if self.remap is not None:
      __magic_name__				     =       indices.reshape(shape[0]			,							-1    )  # add batch axis
      __magic_name__				     =       self.unmap_to_all(UpperCamelCase__    )
      __magic_name__				     =       indices.reshape(-1    )  # flatten again

    # get quantized latent vectors
    __magic_name__				     =       self.embedding(UpperCamelCase__    )

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

    return z_q







class        UpperCAmelCase_			( _A     ):


  '''simple docstring'''







  def __init__( self			: Dict			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: Tuple=False    )			->	Optional[int]:







    """simple docstring"""






    __magic_name__				     =       parameters
    __magic_name__     ,    __magic_name__				     =       torch.chunk(UpperCamelCase__			,							2			,							dim=1    )
    __magic_name__				     =       torch.clamp(self.logvar			,							-30.0			,							20.0    )
    __magic_name__				     =       deterministic
    __magic_name__				     =       torch.exp(0.5 * self.logvar    )
    __magic_name__				     =       torch.exp(self.logvar    )
    if self.deterministic:
      __magic_name__				     =       __magic_name__				     =       torch.zeros_like(
          self.mean			,							device=self.parameters.device			,							dtype=self.parameters.dtype    )
  def 							_lowercase							( self			: Tuple			,							UpperCamelCase__			: Optional[torch.Generator] = None    )			->	torch.FloatTensor:







    """simple docstring"""






    __magic_name__				     =       randn_tensor(
        self.mean.shape			,							generator=UpperCamelCase__			,							device=self.parameters.device			,							dtype=self.parameters.dtype    )
    __magic_name__				     =       self.mean + self.std * sample
    return x
  def 							_lowercase							( self			: Dict			,							UpperCamelCase__			: Optional[int]=None    )			->	Any:







    """simple docstring"""






    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 							_lowercase							( self			: Optional[int]			,							UpperCamelCase__			: List[Any]			,							UpperCamelCase__			: Dict=[1, 2, 3]    )			->	Optional[int]:







    """simple docstring"""






    if self.deterministic:
      return torch.Tensor([0.0]    )
    __magic_name__				     =       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 							_lowercase							( self			: Optional[int]    )			->	Optional[int]:







    """simple docstring"""






    return self.mean

