code: 
  
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  
  
  import functools
  import math
  
  import flax.linen as nn
  import jax
  import jax.numpy as jnp
  
  
  
  
  
  def      _lowerCAmelCase		(		lowercase      ,			lowercase      ,			lowercase      ,			lowercase      ,			lowercase = 4096			)      ->      Union[str, Any]:
     __lowerCAmelCase      		=       key.shape[-3:]
     __lowerCAmelCase      		=       value.shape[-1]
     __lowerCAmelCase      		=       min(_SCREAMING_SNAKE_CASE      ,			_SCREAMING_SNAKE_CASE			)
     __lowerCAmelCase      		=       query / jnp.sqrt(_SCREAMING_SNAKE_CASE			)
  
     @functools.partial(jax.checkpoint      ,			prevent_cse=_SCREAMING_SNAKE_CASE			)
     def summarize_chunk(lowercase      ,			lowercase      ,			lowercase			):
        __lowerCAmelCase      		=       jnp.einsum("""...qhd,...khd->...qhk"""      ,			_SCREAMING_SNAKE_CASE      ,			_SCREAMING_SNAKE_CASE      ,			precision=_SCREAMING_SNAKE_CASE			)
  
        __lowerCAmelCase      		=       jnp.max(_SCREAMING_SNAKE_CASE      ,			axis=-1      ,			keepdims=_SCREAMING_SNAKE_CASE			)
        __lowerCAmelCase      		=       jax.lax.stop_gradient(_SCREAMING_SNAKE_CASE			)
        __lowerCAmelCase      		=       jnp.exp(attn_weights - max_score			)
  
        __lowerCAmelCase      		=       jnp.einsum("""...vhf,...qhv->...qhf"""      ,			_SCREAMING_SNAKE_CASE      ,			_SCREAMING_SNAKE_CASE      ,			precision=_SCREAMING_SNAKE_CASE			)
        __lowerCAmelCase      		=       jnp.einsum("""...qhk->...qh"""      ,			_SCREAMING_SNAKE_CASE			)
  
        return (exp_values, exp_weights.sum(axis=-1			), max_score)
  
     def chunk_scanner(lowercase			):
        # julienne key array
        __lowerCAmelCase      		=       jax.lax.dynamic_slice(
            operand=_SCREAMING_SNAKE_CASE      ,			start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0]      ,			slice_sizes=list(key.shape[:-3]			) + [key_chunk_size, num_heads, k_features]      ,			)
  
        # julienne value array
        __lowerCAmelCase      		=       jax.lax.dynamic_slice(
            operand=_SCREAMING_SNAKE_CASE      ,			start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0]      ,			slice_sizes=list(value.shape[:-3]			) + [key_chunk_size, num_heads, v_features]      ,			)
  
        return summarize_chunk(_SCREAMING_SNAKE_CASE      ,			_SCREAMING_SNAKE_CASE      ,			_SCREAMING_SNAKE_CASE			)
  
     __lowerCAmelCase      		=       jax.lax.map(f=_SCREAMING_SNAKE_CASE      ,			xs=jnp.arange(0      ,			_SCREAMING_SNAKE_CASE      ,			_SCREAMING_SNAKE_CASE			)			)
  
     __lowerCAmelCase      		=       jnp.max(_SCREAMING_SNAKE_CASE      ,			axis=0      ,			keepdims=_SCREAMING_SNAKE_CASE			)
     __lowerCAmelCase      		=       jnp.exp(chunk_max - global_max			)
  
     chunk_values *= jnp.expand_dims(_SCREAMING_SNAKE_CASE      ,			axis=-1			)
     chunk_weights *= max_diffs
  
     __lowerCAmelCase      		=       chunk_values.sum(axis=0			)
     __lowerCAmelCase      		=       jnp.expand_dims(_SCREAMING_SNAKE_CASE      ,			-1			).sum(axis=0			)
  
     return all_values / all_weights
  
  
  
  
  
  def      _lowerCAmelCase		(		lowercase      ,			lowercase      ,			lowercase      ,			lowercase=jax.lax.Precision.HIGHEST      ,			lowercase = 1024      ,			lowercase = 4096			)      ->      Optional[Any]:
     __lowerCAmelCase      		=       query.shape[-3:]
  
     def chunk_scanner(lowercase      ,			lowercase			):
        # julienne query array
        __lowerCAmelCase      		=       jax.lax.dynamic_slice(
            operand=_SCREAMING_SNAKE_CASE      ,			start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]      ,			slice_sizes=list(query.shape[:-3]			) + [min(_SCREAMING_SNAKE_CASE      ,			_SCREAMING_SNAKE_CASE			), num_heads, q_features]      ,			)
  
        return (
            chunk_idx + query_chunk_size,  # unused ignore it
            _query_chunk_attention(
                query=_SCREAMING_SNAKE_CASE      ,			key=_SCREAMING_SNAKE_CASE      ,			value=_SCREAMING_SNAKE_CASE      ,			precision=_SCREAMING_SNAKE_CASE      ,			key_chunk_size=_SCREAMING_SNAKE_CASE			),
        )
  
     __lowerCAmelCase      		=       jax.lax.scan(
         f=_SCREAMING_SNAKE_CASE      ,			init=0      ,			xs=_SCREAMING_SNAKE_CASE      ,			length=math.ceil(num_q / query_chunk_size			)			)
  
     return jnp.concatenate(_SCREAMING_SNAKE_CASE      ,			axis=-3			)  # fuse the chunked result back
  
  
  
  
  
  class 					_UpperCAmelCase							(		nn.Module     ):
    a     :     Any     						=42
    a     :     Any     						=8
    a     :     Union[str, Any]     						=64
    a     :     Optional[int]     						=0.0
    a     :     List[str]     						=False
    a     :     str     						=jnp.floataa
    def 						lowerCamelCase__    (						self    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       self.dim_head * self.heads
       __lowerCAmelCase      		=       self.dim_head**-0.5
  
       # Weights were exported with old names {to_q, to_k, to_v, to_out}
       __lowerCAmelCase      		=       nn.Dense(__a,use_bias=__a,dtype=self.dtype,name="""to_q"""    )
       __lowerCAmelCase      		=       nn.Dense(__a,use_bias=__a,dtype=self.dtype,name="""to_k"""    )
       __lowerCAmelCase      		=       nn.Dense(__a,use_bias=__a,dtype=self.dtype,name="""to_v"""    )
  
       __lowerCAmelCase      		=       nn.Dense(self.query_dim,dtype=self.dtype,name="""to_out_0"""    )
       __lowerCAmelCase      		=       nn.Dropout(rate=self.dropout    )
    def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       tensor.shape
       __lowerCAmelCase      		=       self.heads
       __lowerCAmelCase      		=       tensor.reshape(__a,__a,__a,dim // head_size    )
       __lowerCAmelCase      		=       jnp.transpose(__a,(0, 2, 1, 3)    )
       __lowerCAmelCase      		=       tensor.reshape(batch_size * head_size,__a,dim // head_size    )
       return tensor
    def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       tensor.shape
       __lowerCAmelCase      		=       self.heads
       __lowerCAmelCase      		=       tensor.reshape(batch_size // head_size,__a,__a,__a    )
       __lowerCAmelCase      		=       jnp.transpose(__a,(0, 2, 1, 3)    )
       __lowerCAmelCase      		=       tensor.reshape(batch_size // head_size,__a,dim * head_size    )
       return tensor
  
  
  
  
  
  
  
    def __call__(						self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE=None,__SCREAMING_SNAKE_CASE=True    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       hidden_states if context is None else context
  
       __lowerCAmelCase      		=       self.query(__a    )
       __lowerCAmelCase      		=       self.key(__a    )
       __lowerCAmelCase      		=       self.value(__a    )
  
       __lowerCAmelCase      		=       self.reshape_heads_to_batch_dim(__a    )
       __lowerCAmelCase      		=       self.reshape_heads_to_batch_dim(__a    )
       __lowerCAmelCase      		=       self.reshape_heads_to_batch_dim(__a    )
  
       if self.use_memory_efficient_attention:
          __lowerCAmelCase      		=       query_states.transpose(1,0,2    )
          __lowerCAmelCase      		=       key_states.transpose(1,0,2    )
          __lowerCAmelCase      		=       value_states.transpose(1,0,2    )
  
          # this if statement create a chunk size for each layer of the unet
          # the chunk size is equal to the query_length dimension of the deepest layer of the unet
  
          __lowerCAmelCase      		=       query_states.shape[-3]
          if flatten_latent_dim % 64 == 0:
             __lowerCAmelCase      		=       int(flatten_latent_dim / 64    )
          elif flatten_latent_dim % 16 == 0:
             __lowerCAmelCase      		=       int(flatten_latent_dim / 16    )
          elif flatten_latent_dim % 4 == 0:
             __lowerCAmelCase      		=       int(flatten_latent_dim / 4    )
          else:
             __lowerCAmelCase      		=       int(__a    )
  
          __lowerCAmelCase      		=       jax_memory_efficient_attention(
              __a,__a,__a,query_chunk_size=__a,key_chunk_size=40_96 * 4    )
  
          __lowerCAmelCase      		=       hidden_states.transpose(1,0,2    )
       else:
          # compute attentions
          __lowerCAmelCase      		=       jnp.einsum("""b i d, b j d->b i j""",__a,__a    )
          __lowerCAmelCase      		=       attention_scores * self.scale
          __lowerCAmelCase      		=       nn.softmax(__a,axis=2    )
  
          # attend to values
          __lowerCAmelCase      		=       jnp.einsum("""b i j, b j d -> b i d""",__a,__a    )
  
       __lowerCAmelCase      		=       self.reshape_batch_dim_to_heads(__a    )
       __lowerCAmelCase      		=       self.proj_attn(__a    )
       return self.dropout_layer(__a,deterministic=__a    )
  
  
  
  
  
  class 					_UpperCAmelCase							(		nn.Module     ):
    a     :     Union[str, Any]     						=42
    a     :     Optional[int]     						=42
    a     :     Dict     						=42
    a     :     List[str]     						=0.0
    a     :     Optional[int]     						=False
    a     :     Optional[int]     						=jnp.floataa
    a     :     Optional[Any]     						=False
    def 						lowerCamelCase__    (						self    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       FlaxAttention(
           self.dim,self.n_heads,self.d_head,self.dropout,self.use_memory_efficient_attention,dtype=self.dtype    )
       # cross attention
       __lowerCAmelCase      		=       FlaxAttention(
           self.dim,self.n_heads,self.d_head,self.dropout,self.use_memory_efficient_attention,dtype=self.dtype    )
       __lowerCAmelCase      		=       FlaxFeedForward(dim=self.dim,dropout=self.dropout,dtype=self.dtype    )
       __lowerCAmelCase      		=       nn.LayerNorm(epsilon=1e-5,dtype=self.dtype    )
       __lowerCAmelCase      		=       nn.LayerNorm(epsilon=1e-5,dtype=self.dtype    )
       __lowerCAmelCase      		=       nn.LayerNorm(epsilon=1e-5,dtype=self.dtype    )
       __lowerCAmelCase      		=       nn.Dropout(rate=self.dropout    )
  
  
  
  
  
  
  
    def __call__(						self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE=True    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       hidden_states
       if self.only_cross_attention:
          __lowerCAmelCase      		=       self.attna(self.norma(__a    ),__a,deterministic=__a    )
       else:
          __lowerCAmelCase      		=       self.attna(self.norma(__a    ),deterministic=__a    )
       __lowerCAmelCase      		=       hidden_states + residual
  
       # cross attention
       __lowerCAmelCase      		=       hidden_states
       __lowerCAmelCase      		=       self.attna(self.norma(__a    ),__a,deterministic=__a    )
       __lowerCAmelCase      		=       hidden_states + residual
  
       # feed forward
       __lowerCAmelCase      		=       hidden_states
       __lowerCAmelCase      		=       self.ff(self.norma(__a    ),deterministic=__a    )
       __lowerCAmelCase      		=       hidden_states + residual
  
       return self.dropout_layer(__a,deterministic=__a    )
  
  
  
  
  
  class 					_UpperCAmelCase							(		nn.Module     ):
    a     :     Tuple     						=42
    a     :     Tuple     						=42
    a     :     Optional[Any]     						=42
    a     :     Tuple     						=1
    a     :     Tuple     						=0.0
    a     :     List[Any]     						=False
    a     :     Union[str, Any]     						=False
    a     :     Optional[int]     						=jnp.floataa
    a     :     Union[str, Any]     						=False
    def 						lowerCamelCase__    (						self    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       nn.GroupNorm(num_groups=32,epsilon=1e-5    )
  
       __lowerCAmelCase      		=       self.n_heads * self.d_head
       if self.use_linear_projection:
          __lowerCAmelCase      		=       nn.Dense(__a,dtype=self.dtype    )
       else:
          __lowerCAmelCase      		=       nn.Conv(
              __a,kernel_size=(1, 1),strides=(1, 1),padding="""VALID""",dtype=self.dtype,)
  
       __lowerCAmelCase      		=       [
           FlaxBasicTransformerBlock(
               __a,self.n_heads,self.d_head,dropout=self.dropout,only_cross_attention=self.only_cross_attention,dtype=self.dtype,use_memory_efficient_attention=self.use_memory_efficient_attention,)
           for _ in range(self.depth    )
       ]
  
       if self.use_linear_projection:
          __lowerCAmelCase      		=       nn.Dense(__a,dtype=self.dtype    )
       else:
          __lowerCAmelCase      		=       nn.Conv(
              __a,kernel_size=(1, 1),strides=(1, 1),padding="""VALID""",dtype=self.dtype,)
  
       __lowerCAmelCase      		=       nn.Dropout(rate=self.dropout    )
  
  
  
  
  
  
  
    def __call__(						self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE=True    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       hidden_states.shape
       __lowerCAmelCase      		=       hidden_states
       __lowerCAmelCase      		=       self.norm(__a    )
       if self.use_linear_projection:
          __lowerCAmelCase      		=       hidden_states.reshape(__a,height * width,__a    )
          __lowerCAmelCase      		=       self.proj_in(__a    )
       else:
          __lowerCAmelCase      		=       self.proj_in(__a    )
          __lowerCAmelCase      		=       hidden_states.reshape(__a,height * width,__a    )
  
       for transformer_block in self.transformer_blocks:
          __lowerCAmelCase      		=       transformer_block(__a,__a,deterministic=__a    )
  
       if self.use_linear_projection:
          __lowerCAmelCase      		=       self.proj_out(__a    )
          __lowerCAmelCase      		=       hidden_states.reshape(__a,__a,__a,__a    )
       else:
          __lowerCAmelCase      		=       hidden_states.reshape(__a,__a,__a,__a    )
          __lowerCAmelCase      		=       self.proj_out(__a    )
  
       __lowerCAmelCase      		=       hidden_states + residual
       return self.dropout_layer(__a,deterministic=__a    )
  
  
  
  
  
  class 					_UpperCAmelCase							(		nn.Module     ):
    a     :     Any     						=42
    a     :     Any     						=0.0
    a     :     str     						=jnp.floataa
    def 						lowerCamelCase__    (						self    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       FlaxGEGLU(self.dim,self.dropout,self.dtype    )
       __lowerCAmelCase      		=       nn.Dense(self.dim,dtype=self.dtype    )
  
  
  
  
  
  
  
    def __call__(						self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE=True    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       self.net_a(__a,deterministic=__a    )
       __lowerCAmelCase      		=       self.net_a(__a    )
       return hidden_states
  
  
  
  
  
  class 					_UpperCAmelCase							(		nn.Module     ):
    a     :     Any     						=42
    a     :     Any     						=0.0
    a     :     Union[str, Any]     						=jnp.floataa
    def 						lowerCamelCase__    (						self    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       self.dim * 4
       __lowerCAmelCase      		=       nn.Dense(inner_dim * 2,dtype=self.dtype    )
       __lowerCAmelCase      		=       nn.Dropout(rate=self.dropout    )
  
  
  
  
  
  
  
    def __call__(						self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE=True    ):
  
  
  
       '''simple docstring'''
  
       __lowerCAmelCase      		=       self.proj(__a    )
       __lowerCAmelCase      		=       jnp.split(__a,2,axis=2    )
       return self.dropout_layer(hidden_linear * nn.gelu(__a    ),deterministic=__a    )
  
  
style_context: 
  
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  
  
  import math
  import sys
  
  
  
  
  
  class 					_UpperCAmelCase							:
  		def __init__(						self    ):
  
  
  
  					'''simple docstring'''
  
  					__lowerCAmelCase      		=       0
  					__lowerCAmelCase      		=       []
  					__lowerCAmelCase      		=       {}  # To store the pos of node in array
  		def 						lowerCamelCase__    (						self    ):
  
  
  
  					'''simple docstring'''
  
  					return self.cur_size == 0
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					__lowerCAmelCase      		=       self.left(__SCREAMING_SNAKE_CASE    )
  					__lowerCAmelCase      		=       self.right(__SCREAMING_SNAKE_CASE    )
  					if lc < self.cur_size and self.array(__SCREAMING_SNAKE_CASE    )[0] < self.array(__SCREAMING_SNAKE_CASE    )[0]:
  								__lowerCAmelCase      		=       lc
  					else:
  								__lowerCAmelCase      		=       idx
  					if rc < self.cur_size and self.array(__SCREAMING_SNAKE_CASE    )[0] < self.array(__SCREAMING_SNAKE_CASE    )[0]:
  								__lowerCAmelCase      		=       rc
  					if smallest != idx:
  								self.swap(__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE    )
  								self.min_heapify(__SCREAMING_SNAKE_CASE    )
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					__lowerCAmelCase      		=       self.cur_size
  					self.cur_size += 1
  					self.array.append((sys.maxsize, tup[1])    )
  					self.decrease_key((sys.maxsize, tup[1]),tup[0]    )
  		def 						lowerCamelCase__    (						self    ):
  
  
  
  					'''simple docstring'''
  
  					__lowerCAmelCase      		=       self.array[0][1]
  					__lowerCAmelCase      		=       self.array[self.cur_size - 1]
  					self.cur_size -= 1
  					self.min_heapify(1    )
  					del self.pos[min_node]
  					return min_node
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					return 2 * i + 1
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					return 2 * i + 2
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					return math.floor(i / 2    )
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					__lowerCAmelCase      		=       j
  					__lowerCAmelCase      		=       i
  					__lowerCAmelCase      		=       self.array[i]
  					__lowerCAmelCase      		=       self.array[j]
  					__lowerCAmelCase      		=       temp
  
  
  
  
  
  
  
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					__lowerCAmelCase      		=       self.pos[tup[1]]
  					# assuming the new_d is atmost old_d
  					__lowerCAmelCase      		=       (new_d, tup[1])
  					while idx > 0 and self.array[self.par(__SCREAMING_SNAKE_CASE    )][0] > self.array[idx][0]:
  								self.swap(__SCREAMING_SNAKE_CASE,self.par(__SCREAMING_SNAKE_CASE    )    )
  								__lowerCAmelCase      		=       self.par(__SCREAMING_SNAKE_CASE    )
  
  
  
  
  
  class 					_UpperCAmelCase							:
  		def __init__(						self,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					__lowerCAmelCase      		=       {}  # To store graph: u -> (v,w)
  					__lowerCAmelCase      		=       num  # Number of nodes in graph
  					# To store the distance from source vertex
  					__lowerCAmelCase      		=       [0] * self.num_nodes
  					__lowerCAmelCase      		=       [-1] * self.num_nodes  # To store the path
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					if u in self.adjList:
  								self.adjList[u].append((v, w)    )
  					else:
  								__lowerCAmelCase      		=       [(v, w)]
  
  					# Assuming undirected graph
  					if v in self.adjList:
  								self.adjList[v].append((u, w)    )
  					else:
  								__lowerCAmelCase      		=       [(u, w)]
  		def 						lowerCamelCase__    (						self    ):
  
  
  
  					'''simple docstring'''
  
  					for u in self.adjList:
  								print(__SCREAMING_SNAKE_CASE,"""->""",""" -> """.join(str(f'{v}({w})'    ) for v, w in self.adjList[u]    )    )
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					__lowerCAmelCase      		=       [-1] * self.num_nodes
  					# src is the source node
  					__lowerCAmelCase      		=       0
  					__lowerCAmelCase      		=       PriorityQueue()
  					q.insert((0, src)    )  # (dist from src, node)
  					for u in self.adjList:
  								if u != src:
  											__lowerCAmelCase      		=       sys.maxsize  # Infinity
  											__lowerCAmelCase      		=       -1
  
  					while not q.is_empty():
  								__lowerCAmelCase      		=       q.extract_min()  # Returns node with the min dist from source
  								# Update the distance of all the neighbours of u and
  								# if their prev dist was INFINITY then push them in Q
  								for v, w in self.adjList[u]:
  											__lowerCAmelCase      		=       self.dist[u] + w
  											if self.dist[v] > new_dist:
  														if self.dist[v] == sys.maxsize:
  																	q.insert((new_dist, v)    )
  														else:
  																	q.decrease_key((self.dist[v], v),__SCREAMING_SNAKE_CASE    )
  														__lowerCAmelCase      		=       new_dist
  														__lowerCAmelCase      		=       u
  
          # Show the shortest distances from src
  					self.show_distances(__SCREAMING_SNAKE_CASE    )
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					print(f'Distance from node: {src}'    )
  					for u in range(self.num_nodes    ):
  								print(f'Node {u} has distance: {self.dist[u]}'    )
  
  
  
  
  
  
  
  		def 						lowerCamelCase__    (						self,__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE    ):
  
  
  
  					'''simple docstring'''
  
  					__lowerCAmelCase      		=       []
  					__lowerCAmelCase      		=       0
  					__lowerCAmelCase      		=       dest
  					# Backtracking from dest to src
  					while self.par[temp] != -1:
  								path.append(__SCREAMING_SNAKE_CASE    )
  								if temp != src:
  											for v, w in self.adjList[temp]:
  														if v == self.par[temp]:
  																	cost += w
  																	break
  								__lowerCAmelCase      		=       self.par[temp]
  					path.append(__SCREAMING_SNAKE_CASE    )
  					path.reverse()
  
  					print(f'----Path to reach {dest} from {src}----'    )
  					for u in path:
  								print(f'{u}',end=""" """    )
  								if u != dest:
  											print("""-> """,end=""""""    )
  
  					print("""\nTotal cost of path: """,__SCREAMING_SNAKE_CASE    )
  
  
  if __name__ == "__main__":
  	_a      :				str											=					Graph(9)
  	graph.add_edge(0, 1, 4)
  	graph.add_edge(0, 7, 8)
  	graph.add_edge(1, 2, 8)
  	graph.add_edge(1, 7, 1_1)
  	graph.add_edge(2, 3, 7)
  	graph.add_edge(2, 8, 2)
  	graph.add_edge(2, 5, 4)
  	graph.add_edge(3, 4, 9)
  	graph.add_edge(3, 5, 1_4)
  	graph.add_edge(4, 5, 1_0)
  	graph.add_edge(5, 6, 2)
  	graph.add_edge(6, 7, 1)
  	graph.add_edge(6, 8, 6)
  	graph.add_edge(7, 8, 7)
  	graph.show_graph()
  	graph.dijkstra(0)
  	graph.show_path(0, 4)
  
  # OUTPUT
  # 0 -> 1(4) -> 7(8)
  # 1 -> 0(4) -> 2(8) -> 7(11)
  # 7 -> 0(8) -> 1(11) -> 6(1) -> 8(7)
  # 2 -> 1(8) -> 3(7) -> 8(2) -> 5(4)
  # 3 -> 2(7) -> 4(9) -> 5(14)
  # 8 -> 2(2) -> 6(6) -> 7(7)
  # 5 -> 2(4) -> 3(14) -> 4(10) -> 6(2)
  # 4 -> 3(9) -> 5(10)
  # 6 -> 5(2) -> 7(1) -> 8(6)
  # Distance from node: 0
  # Node 0 has distance: 0
  # Node 1 has distance: 4
  # Node 2 has distance: 12
  # Node 3 has distance: 19
  # Node 4 has distance: 21
  # Node 5 has distance: 11
  # Node 6 has distance: 9
  # Node 7 has distance: 8
  # Node 8 has distance: 14
  # ----Path to reach 4 from 0----
  # 0 -> 7 -> 6 -> 5 -> 4
  # Total cost of path:  21
  
  