





'''simple docstring'''

from collections import deque
from math import floor
from random import random
from time import time






class   UpperCAmelCase      :
 def __init__(							self :Optional[Any]				)->				Any:
     A__							=     {}
 def 	UpperCAmelCase_				(							self :Dict		,			lowercase_ :Optional[int]		,			lowercase_ :Dict		,			lowercase_ :List[Any]=1				)->				Any:
     if self.graph.get(lowercase_				):
         if self.graph[u].count([w, v]				) == 0:
             self.graph[u].append([w, v]				)
     else:
         A__							=     [[w, v]]
     if not self.graph.get(lowercase_				):
         A__							=     []
 def 	UpperCAmelCase_				(							self :int				)->				Any:
     return list(self.graph				)
 def 	UpperCAmelCase_				(							self :List[Any]		,			lowercase_ :Union[str, Any]		,			lowercase_ :Union[str, Any]				)->				List[Any]:
     if self.graph.get(lowercase_				):
         for _ in self.graph[u]:
             if _[1] == v:
                 self.graph[u].remove(lowercase_				)
 def 	UpperCAmelCase_				(							self :List[Any]		,			lowercase_ :int=-2		,			lowercase_ :Optional[int]=-1				)->				int:
     if s == d:
         return []
     A__							=     []
     A__							=     []
     if s == -2:
         A__							=     list(self.graph				)[0]
     stack.append(lowercase_				)
     visited.append(lowercase_				)
     A__							=     s

     while True:
         # check if there is any non isolated nodes
         if len(self.graph[s]				) != 0:
             A__							=     s
             for node in self.graph[s]:
                 if visited.count(node[1]				) < 1:
                     if node[1] == d:
                         visited.append(lowercase_				)
                         return visited
                     else:
                         stack.append(node[1]				)
                         visited.append(node[1]				)
                         A__							=     node[1]
                         break

            # check if all the children are visited
         if s == ss:
             stack.pop()
             if len(lowercase_				) != 0:
                 A__							=     stack[len(lowercase_				) - 1]
         else:
             A__							=     ss

         # check if se have reached the starting point
         if len(lowercase_				) == 0:
             return visited
 def 	UpperCAmelCase_				(							self :Dict		,			lowercase_ :Optional[int]=-1				)->				Optional[int]:
     if c == -1:
         A__							=     floor(random() * 1_00_00				) + 10
     for i in range(lowercase_				):
         # every vertex has max 100 edges
         for _ in range(floor(random() * 1_02				) + 1				):
             A__							=     floor(random() * c				) + 1
             if n != i:
                 self.add_pair(lowercase_		,			lowercase_		,			1				)
 def 	UpperCAmelCase_				(							self :Dict		,			lowercase_ :Optional[Any]=-2				)->				int:
     A__							=     deque()
     A__							=     []
     if s == -2:
         A__							=     list(self.graph				)[0]
     d.append(lowercase_				)
     visited.append(lowercase_				)
     while d:
         A__							=     d.popleft()
         if len(self.graph[s]				) != 0:
             for node in self.graph[s]:
                 if visited.count(node[1]				) < 1:
                     d.append(node[1]				)
                     visited.append(node[1]				)
     return visited
 def 	UpperCAmelCase_				(							self :Dict		,			lowercase_ :Optional[int]				)->				Union[str, Any]:
     A__							=     0
     for x in self.graph:
         for y in self.graph[x]:
             if y[1] == u:
                 count += 1
     return count
 def 	UpperCAmelCase_				(							self :int		,			lowercase_ :Dict				)->				Optional[int]:
     return len(self.graph[u]				)
 def 	UpperCAmelCase_				(							self :Tuple		,			lowercase_ :Optional[int]=-2				)->				Union[str, Any]:
     A__							=     []
     A__							=     []
     if s == -2:
         A__							=     list(self.graph				)[0]
     stack.append(lowercase_				)
     visited.append(lowercase_				)
     A__							=     s
     A__							=     []

     while True:
         # check if there is any non isolated nodes
         if len(self.graph[s]				) != 0:
             A__							=     s
             for node in self.graph[s]:
                 if visited.count(node[1]				) < 1:
                     stack.append(node[1]				)
                     visited.append(node[1]				)
                     A__							=     node[1]
                     break

            # check if all the children are visited
         if s == ss:
             sorted_nodes.append(stack.pop()				)
             if len(lowercase_				) != 0:
                 A__							=     stack[len(lowercase_				) - 1]
         else:
             A__							=     ss

         # check if se have reached the starting point
         if len(lowercase_				) == 0:
             return sorted_nodes
 def 	UpperCAmelCase_				(							self :Optional[Any]				)->				List[Any]:
     A__							=     []
     A__							=     []
     A__							=     list(self.graph				)[0]
     stack.append(lowercase_				)
     visited.append(lowercase_				)
     A__							=     -2
     A__							=     []
     A__							=     s
     A__							=     False
     A__							=     set()

     while True:
         # check if there is any non isolated nodes
         if len(self.graph[s]				) != 0:
             A__							=     s
             for node in self.graph[s]:
                 if (
                     visited.count(node[1]				) > 0
                     and node[1] != parent
                     and indirect_parents.count(node[1]				) > 0
                     and not on_the_way_back
                 ):
                     A__							=     len(lowercase_				) - 1
                     while len_stack >= 0:
                         if stack[len_stack] == node[1]:
                             anticipating_nodes.add(node[1]				)
                             break
                         else:
                             anticipating_nodes.add(stack[len_stack]				)
                             len_stack -= 1
                 if visited.count(node[1]				) < 1:
                     stack.append(node[1]				)
                     visited.append(node[1]				)
                     A__							=     node[1]
                     break

            # check if all the children are visited
         if s == ss:
             stack.pop()
             A__							=     True
             if len(lowercase_				) != 0:
                 A__							=     stack[len(lowercase_				) - 1]
         else:
             A__							=     False
             indirect_parents.append(lowercase_				)
             A__							=     s
             A__							=     ss

         # check if se have reached the starting point
         if len(lowercase_				) == 0:
             return list(lowercase_				)
 def 	UpperCAmelCase_				(							self :str				)->				List[str]:
     A__							=     []
     A__							=     []
     A__							=     list(self.graph				)[0]
     stack.append(lowercase_				)
     visited.append(lowercase_				)
     A__							=     -2
     A__							=     []
     A__							=     s
     A__							=     False
     A__							=     set()

     while True:
         # check if there is any non isolated nodes
         if len(self.graph[s]				) != 0:
             A__							=     s
             for node in self.graph[s]:
                 if (
                     visited.count(node[1]				) > 0
                     and node[1] != parent
                     and indirect_parents.count(node[1]				) > 0
                     and not on_the_way_back
                 ):
                     A__							=     len(lowercase_				) - 1
                     while len_stack_minus_one >= 0:
                         if stack[len_stack_minus_one] == node[1]:
                             anticipating_nodes.add(node[1]				)
                             break
                         else:
                             return True
                 if visited.count(node[1]				) < 1:
                     stack.append(node[1]				)
                     visited.append(node[1]				)
                     A__							=     node[1]
                     break

            # check if all the children are visited
         if s == ss:
             stack.pop()
             A__							=     True
             if len(lowercase_				) != 0:
                 A__							=     stack[len(lowercase_				) - 1]
         else:
             A__							=     False
             indirect_parents.append(lowercase_				)
             A__							=     s
             A__							=     ss

         # check if se have reached the starting point
         if len(lowercase_				) == 0:
             return False
 def 	UpperCAmelCase_				(							self :Union[str, Any]		,			lowercase_ :Optional[Any]=-2		,			lowercase_ :Any=-1				)->				List[str]:
     A__							=     time()
     self.dfs(lowercase_		,			lowercase_				)
     A__							=     time()
     return end - begin



 def 	UpperCAmelCase_				(							self :Union[str, Any]		,			lowercase_ :List[str]=-2				)->				Tuple:
     A__							=     time()
     self.bfs(lowercase_				)
     A__							=     time()
     return end - begin






class   UpperCAmelCase      :
 def __init__(							self :Optional[Any]				)->				List[str]:
     A__							=     {}
 def 	UpperCAmelCase_				(							self :Optional[Any]		,			lowercase_ :str		,			lowercase_ :int		,			lowercase_ :Optional[Any]=1				)->				List[str]:
     # check if the u exists
     if self.graph.get(lowercase_				):
         # if there already is a edge
         if self.graph[u].count([w, v]				) == 0:
             self.graph[u].append([w, v]				)
     else:
         # if u does not exist
         A__							=     [[w, v]]
     # add the other way
     if self.graph.get(lowercase_				):
         # if there already is a edge
         if self.graph[v].count([w, u]				) == 0:
             self.graph[v].append([w, u]				)
     else:
         # if u does not exist
         A__							=     [[w, u]]
 def 	UpperCAmelCase_				(							self :str		,			lowercase_ :Dict		,			lowercase_ :List[Any]				)->				Optional[int]:
     if self.graph.get(lowercase_				):
         for _ in self.graph[u]:
             if _[1] == v:
                 self.graph[u].remove(lowercase_				)
        # the other way round
     if self.graph.get(lowercase_				):
         for _ in self.graph[v]:
             if _[1] == u:
                 self.graph[v].remove(lowercase_				)
 def 	UpperCAmelCase_				(							self :List[str]		,			lowercase_ :Optional[int]=-2		,			lowercase_ :str=-1				)->				List[str]:
     if s == d:
         return []
     A__							=     []
     A__							=     []
     if s == -2:
         A__							=     list(self.graph				)[0]
     stack.append(lowercase_				)
     visited.append(lowercase_				)
     A__							=     s

     while True:
         # check if there is any non isolated nodes
         if len(self.graph[s]				) != 0:
             A__							=     s
             for node in self.graph[s]:
                 if visited.count(node[1]				) < 1:
                     if node[1] == d:
                         visited.append(lowercase_				)
                         return visited
                     else:
                         stack.append(node[1]				)
                         visited.append(node[1]				)
                         A__							=     node[1]
                         break

            # check if all the children are visited
         if s == ss:
             stack.pop()
             if len(lowercase_				) != 0:
                 A__							=     stack[len(lowercase_				) - 1]
         else:
             A__							=     ss

         # check if se have reached the starting point
         if len(lowercase_				) == 0:
             return visited
 def 	UpperCAmelCase_				(							self :Any		,			lowercase_ :Union[str, Any]=-1				)->				Union[str, Any]:
     if c == -1:
         A__							=     floor(random() * 1_00_00				) + 10
     for i in range(lowercase_				):
         # every vertex has max 100 edges
         for _ in range(floor(random() * 1_02				) + 1				):
             A__							=     floor(random() * c				) + 1
             if n != i:
                 self.add_pair(lowercase_		,			lowercase_		,			1				)
 def 	UpperCAmelCase_				(							self :Any		,			lowercase_ :int=-2				)->				List[Any]:
     A__							=     deque()
     A__							=     []
     if s == -2:
         A__							=     list(self.graph				)[0]
     d.append(lowercase_				)
     visited.append(lowercase_				)
     while d:
         A__							=     d.popleft()
         if len(self.graph[s]				) != 0:
             for node in self.graph[s]:
                 if visited.count(node[1]				) < 1:
                     d.append(node[1]				)
                     visited.append(node[1]				)
     return visited
 def 	UpperCAmelCase_				(							self :int		,			lowercase_ :Optional[Any]				)->				Optional[Any]:
     return len(self.graph[u]				)
 def 	UpperCAmelCase_				(							self :Optional[Any]				)->				str:
     A__							=     []
     A__							=     []
     A__							=     list(self.graph				)[0]
     stack.append(lowercase_				)
     visited.append(lowercase_				)
     A__							=     -2
     A__							=     []
     A__							=     s
     A__							=     False
     A__							=     set()

     while True:
         # check if there is any non isolated nodes
         if len(self.graph[s]				) != 0:
             A__							=     s
             for node in self.graph[s]:
                 if (
                     visited.count(node[1]				) > 0
                     and node[1] != parent
                     and indirect_parents.count(node[1]				) > 0
                     and not on_the_way_back
                 ):
                     A__							=     len(lowercase_				) - 1
                     while len_stack >= 0:
                         if stack[len_stack] == node[1]:
                             anticipating_nodes.add(node[1]				)
                             break
                         else:
                             anticipating_nodes.add(stack[len_stack]				)
                             len_stack -= 1
                 if visited.count(node[1]				) < 1:
                     stack.append(node[1]				)
                     visited.append(node[1]				)
                     A__							=     node[1]
                     break

            # check if all the children are visited
         if s == ss:
             stack.pop()
             A__							=     True
             if len(lowercase_				) != 0:
                 A__							=     stack[len(lowercase_				) - 1]
         else:
             A__							=     False
             indirect_parents.append(lowercase_				)
             A__							=     s
             A__							=     ss

         # check if se have reached the starting point
         if len(lowercase_				) == 0:
             return list(lowercase_				)
 def 	UpperCAmelCase_				(							self :Optional[int]				)->				Union[str, Any]:
     A__							=     []
     A__							=     []
     A__							=     list(self.graph				)[0]
     stack.append(lowercase_				)
     visited.append(lowercase_				)
     A__							=     -2
     A__							=     []
     A__							=     s
     A__							=     False
     A__							=     set()

     while True:
         # check if there is any non isolated nodes
         if len(self.graph[s]				) != 0:
             A__							=     s
             for node in self.graph[s]:
                 if (
                     visited.count(node[1]				) > 0
                     and node[1] != parent
                     and indirect_parents.count(node[1]				) > 0
                     and not on_the_way_back
                 ):
                     A__							=     len(lowercase_				) - 1
                     while len_stack_minus_one >= 0:
                         if stack[len_stack_minus_one] == node[1]:
                             anticipating_nodes.add(node[1]				)
                             break
                         else:
                             return True
                 if visited.count(node[1]				) < 1:
                     stack.append(node[1]				)
                     visited.append(node[1]				)
                     A__							=     node[1]
                     break

            # check if all the children are visited
         if s == ss:
             stack.pop()
             A__							=     True
             if len(lowercase_				) != 0:
                 A__							=     stack[len(lowercase_				) - 1]
         else:
             A__							=     False
             indirect_parents.append(lowercase_				)
             A__							=     s
             A__							=     ss

         # check if se have reached the starting point
         if len(lowercase_				) == 0:
             return False
 def 	UpperCAmelCase_				(							self :int				)->				Union[str, Any]:
     return list(self.graph				)
 def 	UpperCAmelCase_				(							self :Optional[Any]		,			lowercase_ :List[str]=-2		,			lowercase_ :str=-1				)->				int:
     A__							=     time()
     self.dfs(lowercase_		,			lowercase_				)
     A__							=     time()
     return end - begin



 def 	UpperCAmelCase_				(							self :int		,			lowercase_ :List[str]=-2				)->				int:
     A__							=     time()
     self.bfs(lowercase_				)
     A__							=     time()
     return end - begin










'''simple docstring'''

from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer






@dataclass
class   UpperCAmelCase      (			UpperCamelCase__							):
 __lowercase  						=							42






class   UpperCAmelCase      (			UpperCamelCase__     ,			UpperCamelCase__							):
 @register_to_config
 def __init__(							self :Tuple		,			lowercase_ :int = 3		,			lowercase_ :int = 3		,			lowercase_ :Tuple[str] = ("DownEncoderBlock2D",)		,			lowercase_ :Tuple[str] = ("UpDecoderBlock2D",)		,			lowercase_ :Tuple[int] = (64,)		,			lowercase_ :int = 1		,			lowercase_ :str = "silu"		,			lowercase_ :int = 3		,			lowercase_ :int = 32		,			lowercase_ :int = 2_56		,			lowercase_ :int = 32		,			lowercase_ :Optional[int] = None		,			lowercase_ :float = 0.1_8_2_1_5		,			lowercase_ :str = "group"		,			)->				Optional[int]:
     super().__init__()

     # pass init params to Encoder
     A__							=     Encoder(
         in_channels=lowercase_		,			out_channels=lowercase_		,			down_block_types=lowercase_		,			block_out_channels=lowercase_		,			layers_per_block=lowercase_		,			act_fn=lowercase_		,			norm_num_groups=lowercase_		,			double_z=lowercase_		,			)

     A__							=     vq_embed_dim if vq_embed_dim is not None else latent_channels

     A__							=     nn.Convad(lowercase_		,			lowercase_		,			1				)
     A__							=     VectorQuantizer(lowercase_		,			lowercase_		,			beta=0.2_5		,			remap=lowercase_		,			sane_index_shape=lowercase_				)
     A__							=     nn.Convad(lowercase_		,			lowercase_		,			1				)

     # pass init params to Decoder
     A__							=     Decoder(
         in_channels=lowercase_		,			out_channels=lowercase_		,			up_block_types=lowercase_		,			block_out_channels=lowercase_		,			layers_per_block=lowercase_		,			act_fn=lowercase_		,			norm_num_groups=lowercase_		,			norm_type=lowercase_		,			)
 @apply_forward_hook
 def 	UpperCAmelCase_				(							self :str		,			lowercase_ :torch.FloatTensor		,			lowercase_ :bool = True				)->				VQEncoderOutput:
     A__							=     self.encoder(lowercase_				)
     A__							=     self.quant_conv(lowercase_				)

     if not return_dict:
         return (h,)

     return VQEncoderOutput(latents=lowercase_				)
 @apply_forward_hook
 def 	UpperCAmelCase_				(							self :Dict		,			lowercase_ :torch.FloatTensor		,			lowercase_ :bool = False		,			lowercase_ :bool = True				)->				Union[DecoderOutput, torch.FloatTensor]:
     # also go through quantization layer
     if not force_not_quantize:
         A__,      A__,      A__							=     self.quantize(lowercase_				)
     else:
         A__							=     h
     A__							=     self.post_quant_conv(lowercase_				)
     A__							=     self.decoder(lowercase_		,			quant if self.config.norm_type == "spatial" else None				)

     if not return_dict:
         return (dec,)

     return DecoderOutput(sample=lowercase_				)



 def 	UpperCAmelCase_				(							self :str		,			lowercase_ :torch.FloatTensor		,			lowercase_ :bool = True				)->				Union[DecoderOutput, torch.FloatTensor]:
     A__							=     sample
     A__							=     self.encode(lowercase_				).latents
     A__							=     self.decode(lowercase_				).sample

     if not return_dict:
         return (dec,)

     return DecoderOutput(sample=lowercase_				)



