


'''simple docstring'''



from __future__ import annotations

from sys import maxsize
from typing import Generic, TypeVar

A__						:   Optional[Any]				=    TypeVar("""T""")



def      UpperCAmelCase__   (      UpperCAmelCase_       :							int				)	->						int:
     return (position - 1) // 2



def      UpperCAmelCase__   (      UpperCAmelCase_       :							int				)	->						int:
     return (2 * position) + 1



def      UpperCAmelCase__   (      UpperCAmelCase_       :							int				)	->						int:
     return (2 * position) + 2





class   UpperCAmelCase_    (Generic[T]    ):





    """simple docstring"""






    def __init__(				self      )							->	None:
         __lowerCamelCase  :      list[tuple[T, int]]														=						[]
         __lowerCamelCase  :      dict[T, int]														=						{}
         __lowerCamelCase  :      int														=						0







    def __len__(				self      )							->	int:
         return self.elements







    def __repr__(				self      )							->	str:
         return str(self.heap      )







    def  lowercase_							(				self      )							->	bool:
         # Check if the priority queue is empty
         return self.elements == 0







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )							->	None:
         # Add an element with given priority to the queue
         self.heap.append((elem, weight)      )
         __lowerCamelCase  :      Any														=						self.elements
         self.elements += 1
         self._bubble_up(SCREAMING_SNAKE_CASE_      )







    def  lowercase_							(				self      )							->	T:
         # Remove and return the element with lowest weight (highest priority)
         if self.elements > 1:
              self._swap_nodes(0     ,	self.elements - 1      )
         __lowerCamelCase							,			__lowerCamelCase  :      int														=						self.heap.pop()
         del self.position_map[elem]
         self.elements -= 1
         if self.elements > 0:
              __lowerCamelCase							,			__lowerCamelCase  :      Optional[int]														=						self.heap[0]
              self._bubble_down(SCREAMING_SNAKE_CASE_      )
         return elem







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )							->	None:
         # Update the weight of the given key
         __lowerCamelCase  :      Union[str, Any]														=						self.position_map[elem]
         __lowerCamelCase  :      Any														=						(elem, weight)
         if position > 0:
              __lowerCamelCase  :      Union[str, Any]														=						get_parent_position(SCREAMING_SNAKE_CASE_      )
              __lowerCamelCase							,			__lowerCamelCase  :      List[str]														=						self.heap[parent_position]
              if parent_weight > weight:
                   self._bubble_up(SCREAMING_SNAKE_CASE_      )
              else:
                   self._bubble_down(SCREAMING_SNAKE_CASE_      )
         else:
              self._bubble_down(SCREAMING_SNAKE_CASE_      )







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_      )							->	None:
         # Place a node at the proper position (upward movement) [to be used internally
         # only]
         __lowerCamelCase  :      Tuple														=						self.position_map[elem]
         if curr_pos == 0:
              return None
         __lowerCamelCase  :      Dict														=						get_parent_position(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase							,			__lowerCamelCase  :      int														=						self.heap[curr_pos]
         __lowerCamelCase							,			__lowerCamelCase  :      Union[str, Any]														=						self.heap[parent_position]
         if parent_weight > weight:
              self._swap_nodes(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )
              return self._bubble_up(SCREAMING_SNAKE_CASE_      )
         return None







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_      )							->	None:
         # Place a node at the proper position (downward movement) [to be used
         # internally only]
         __lowerCamelCase  :      List[str]														=						self.position_map[elem]
         __lowerCamelCase							,			__lowerCamelCase  :      int														=						self.heap[curr_pos]
         __lowerCamelCase  :      Union[str, Any]														=						get_child_left_position(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      List[Any]														=						get_child_right_position(SCREAMING_SNAKE_CASE_      )
         if child_left_position < self.elements and child_right_position < self.elements:
              __lowerCamelCase							,			__lowerCamelCase  :      Dict														=						self.heap[child_left_position]
              __lowerCamelCase							,			__lowerCamelCase  :      List[Any]														=						self.heap[child_right_position]
              if child_right_weight < child_left_weight and child_right_weight < weight:
                   self._swap_nodes(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )
                   return self._bubble_down(SCREAMING_SNAKE_CASE_      )
         if child_left_position < self.elements:
              __lowerCamelCase							,			__lowerCamelCase  :      Optional[int]														=						self.heap[child_left_position]
              if child_left_weight < weight:
                   self._swap_nodes(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )
                   return self._bubble_down(SCREAMING_SNAKE_CASE_      )
         else:
              return None
         if child_right_position < self.elements:
              __lowerCamelCase							,			__lowerCamelCase  :      List[str]														=						self.heap[child_right_position]
              if child_right_weight < weight:
                   self._swap_nodes(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )
                   return self._bubble_down(SCREAMING_SNAKE_CASE_      )
         return None







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )							->	None:
         # Swap the nodes at the given positions
         __lowerCamelCase  :      List[Any]														=						self.heap[nodea_pos][0]
         __lowerCamelCase  :      Tuple														=						self.heap[nodea_pos][0]
         __lowerCamelCase							,			__lowerCamelCase  :      Optional[Any]														=						(
             self.heap[nodea_pos],
             self.heap[nodea_pos],
         )
         __lowerCamelCase  :      int														=						nodea_pos
         __lowerCamelCase  :      Optional[Any]														=						nodea_pos







class   UpperCAmelCase_    (Generic[T]    ):





    """simple docstring"""






    def __init__(				self      )							->	None:
         __lowerCamelCase  :      dict[T, dict[T, int]]														=						{}
         __lowerCamelCase  :      int														=						0







    def __repr__(				self      )							->	str:
         return str(self.connections      )







    def __len__(				self      )							->	int:
         return self.nodes







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_      )							->	None:
         # Add a node in the graph if it is not in the graph
         if node not in self.connections:
              __lowerCamelCase  :      Optional[Any]														=						{}
              self.nodes += 1







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      )							->	None:
         # Add an edge between 2 nodes in the graph
         self.add_node(SCREAMING_SNAKE_CASE_      )
         self.add_node(SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      int														=						weight
         __lowerCamelCase  :      Dict														=						weight







def      UpperCAmelCase__   (      UpperCAmelCase_       :							GraphUndirectedWeighted[T] ,       )	->						tuple[dict[T, int], dict[T, T | None]]:
     __lowerCamelCase  :      dict[T, int]														=						{node: maxsize for node in graph.connections}
     __lowerCamelCase  :      dict[T, T | None]														=						{node: None for node in graph.connections}

     __lowerCamelCase  :      MinPriorityQueue[T]														=						MinPriorityQueue()
     for node, weight in dist.items():
          priority_queue.push(UpperCAmelCase_ ,       UpperCAmelCase_				)

     if priority_queue.is_empty():
          return dist, parent

     # initialization
     __lowerCamelCase  :      List[Any]														=						priority_queue.extract_min()
     __lowerCamelCase  :      Union[str, Any]														=						0
     for neighbour in graph.connections[node]:
          if dist[neighbour] > dist[node] + graph.connections[node][neighbour]:
               __lowerCamelCase  :      int														=						dist[node] + graph.connections[node][neighbour]
               priority_queue.update_key(UpperCAmelCase_ ,       dist[neighbour]				)
               __lowerCamelCase  :      int														=						node

    # running prim's algorithm
     while not priority_queue.is_empty():
          __lowerCamelCase  :      List[Any]														=						priority_queue.extract_min()
          for neighbour in graph.connections[node]:
               if dist[neighbour] > dist[node] + graph.connections[node][neighbour]:
                    __lowerCamelCase  :      int														=						dist[node] + graph.connections[node][neighbour]
                    priority_queue.update_key(UpperCAmelCase_ ,       dist[neighbour]				)
                    __lowerCamelCase  :      Union[str, Any]														=						node
     return dist, parent







'''simple docstring'''




import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
__UpperCAmelCase					=object()

# For specifying empty leaf dict `{}`
__UpperCAmelCase					=object()






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							)      ->       List[Any]:
    __lowerCamelCase            =    tuple((re.compile(x + '''$'''							) for x in qs)							)
    for i in range(len(UpperCamelCase__							) - len(UpperCamelCase__							) + 1							):
        __lowerCamelCase            =    [x.match(UpperCamelCase__							) for x, y in zip(UpperCamelCase__							,    ks[i:]							)]
        if matches and all(UpperCamelCase__							):
            return True
    return False






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       Optional[int]:
    def replace(UpperCamelCase__							,    UpperCamelCase__							):
        for rule, replacement in rules:
            if _match(UpperCamelCase__							,    UpperCamelCase__							):
                return replacement
        return val

    return replace






def __lowerCAmelCase					(				)      ->       Dict:
    return [
        # embeddings
        (("transformer", "wpe", "embedding"), P('''mp'''							,    UpperCamelCase__							)),
        (("transformer", "wte", "embedding"), P('''mp'''							,    UpperCamelCase__							)),
        # atention
        (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(UpperCamelCase__							,    '''mp'''							)),
        (("attention", "out_proj", "kernel"), P('''mp'''							,    UpperCamelCase__							)),
        (("attention", "out_proj", "bias"), None),
        # mlp
        (("mlp", "c_fc", "kernel"), P(UpperCamelCase__							,    '''mp'''							)),
        (("mlp", "c_fc", "bias"), P('''mp'''							)),
        (("mlp", "c_proj", "kernel"), P('''mp'''							,    UpperCamelCase__							)),
        (("mlp", "c_proj", "bias"), None),
        # layer norms
        ((r"ln_\d+", "bias"), None),
        ((r"\d+", r"ln_\d+", "scale"), None),
        (("ln_f", "bias"), None),
        (("ln_f", "scale"), None),
    ]






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       int:
    __lowerCamelCase            =    _get_partition_rules()
    __lowerCamelCase            =    _replacement_rules(UpperCamelCase__							)
    __lowerCamelCase            =    {k: _unmatched for k in flatten_dict(UpperCamelCase__							)}
    __lowerCamelCase            =    {k: replace(UpperCamelCase__							,    UpperCamelCase__							) for k, v in initd.items()}
    assert _unmatched not in result.values(), "Incomplete partition spec."
    return freeze(unflatten_dict(UpperCamelCase__							)							)
