

'''simple docstring'''





from __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

UpperCamelCase_    						=       TypeVar("""T""")
UpperCamelCase_    						=       TypeVar("""U""")





class 		__SCREAMING_SNAKE_CASE							(			Generic[T, U]      ):







				def __init__(       self     :		List[str]		,					UpperCAmelCase__     :		T | None		,					UpperCAmelCase__     :		U | None   ):
										'''simple docstring'''
										lowercase :							List[str]    		=key
										lowercase :							Optional[Any]    		=val
										lowercase :							DoubleLinkedListNode[T, U] | None    		=None
										lowercase :							DoubleLinkedListNode[T, U] | None    		=None







				def __repr__(       self     :		Any   ):
										'''simple docstring'''
										return (
										    F'''Node: key: {self.key}, val: {self.val}, '''
										    F'''has next: {bool(self.next   )}, has prev: {bool(self.prev   )}'''
										)





class 		__SCREAMING_SNAKE_CASE							(			Generic[T, U]      ):







				def __init__(       self     :		int   ):
										'''simple docstring'''
										lowercase :							DoubleLinkedListNode[T, U]    		=DoubleLinkedListNode(UpperCAmelCase__		,					UpperCAmelCase__   )
										lowercase :							DoubleLinkedListNode[T, U]    		=DoubleLinkedListNode(UpperCAmelCase__		,					UpperCAmelCase__   )
										lowercase		,	lowercase :							Any    		=self.rear, self.head







				def __repr__(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										lowercase :							Any    		=['''DoubleLinkedList''']
										lowercase :							Union[str, Any]    		=self.head
										while node.next is not None:
																rep.append(str(UpperCAmelCase__   )   )
																lowercase :							Any    		=node.next
										rep.append(str(self.rear   )   )
										return ",\n    ".join(UpperCAmelCase__   )







				def  lowerCamelCase_	(       self     :		int		,					UpperCAmelCase__     :		DoubleLinkedListNode[T, U]   ):
										'''simple docstring'''
										lowercase :							Optional[int]    		=self.rear.prev

										# All nodes other than self.head are guaranteed to have non-None previous
										assert previous is not None

										lowercase :							Dict    		=node
										lowercase :							str    		=previous
										lowercase :							Optional[Any]    		=node
										lowercase :							Optional[int]    		=self.rear







				def  lowerCamelCase_	(       self     :		Optional[int]		,					UpperCAmelCase__     :		DoubleLinkedListNode[T, U]   ):
										'''simple docstring'''

										if node.prev is None or node.next is None:
																return None

										lowercase :							Any    		=node.next
										lowercase :							Optional[int]    		=node.prev
										lowercase :							Any    		=None
										lowercase :							int    		=None
										return node





class 		__SCREAMING_SNAKE_CASE							(			Generic[T, U]      ):
				lowerCamelCase_		       =    {}







				def __init__(       self     :		List[Any]		,					UpperCAmelCase__     :		int   ):
										'''simple docstring'''
										lowercase :							DoubleLinkedList[T, U]    		=DoubleLinkedList()
										lowercase :							Union[str, Any]    		=capacity
										lowercase :							int    		=0
										lowercase :							List[Any]    		=0
										lowercase :							Optional[Any]    		=0
										lowercase :							dict[T, DoubleLinkedListNode[T, U]]    		={}







				def __repr__(       self     :		List[str]   ):
										'''simple docstring'''
										return (
										    F'''CacheInfo(hits={self.hits}, misses={self.miss}, '''
										    F'''capacity={self.capacity}, current size={self.num_keys})'''
										)







				def __contains__(       self     :		Dict		,					UpperCAmelCase__     :		T   ):
										'''simple docstring'''
										return key in self.cache







				def  lowerCamelCase_	(       self     :		Union[str, Any]		,					UpperCAmelCase__     :		T   ):
										'''simple docstring'''
										# Note: pythonic interface would throw KeyError rather than return None

										if key in self.cache:
																self.hits += 1
																lowercase :							DoubleLinkedListNode[T, U]    		=self.cache[key]
																lowercase :							int    		=self.list.remove(self.cache[key]   )
																assert node == value_node

																# node is guaranteed not None because it is in self.cache
																assert node is not None
																self.list.add(UpperCAmelCase__   )
																return node.val
										self.miss += 1
										return None







				def  lowerCamelCase_	(       self     :		Optional[Any]		,					UpperCAmelCase__     :		T		,					UpperCAmelCase__     :		U   ):
										'''simple docstring'''

										if key not in self.cache:
																if self.num_keys >= self.capacity:
																						# delete first node (oldest) when over capacity
																						lowercase :							Dict    		=self.list.head.next

																						# guaranteed to have a non-None first node when num_keys > 0
																						# explain to type checker via assertions
																						assert first_node is not None
																						assert first_node.key is not None
																						assert (
																						    self.list.remove(UpperCAmelCase__   ) is not None
																						)  # node guaranteed to be in list assert node.key is not None

																						del self.cache[first_node.key]
																						self.num_keys -= 1
																lowercase :							Dict    		=DoubleLinkedListNode(UpperCAmelCase__		,					UpperCAmelCase__   )
																self.list.add(self.cache[key]   )
																self.num_keys += 1

										else:
																# bump node to the end of the list, update value
																lowercase :							Union[str, Any]    		=self.list.remove(self.cache[key]   )
																assert node is not None  # node guaranteed to be in list
																lowercase :							Any    		=value
																self.list.add(UpperCAmelCase__   )







				@classmethod
				def  lowerCamelCase_	(       cls     :		Optional[Any]		,					UpperCAmelCase__     :		int = 128   ):
										'''simple docstring'''

										def cache_decorator_inner(UpperCAmelCase__     :		Callable[[T], U]   ) -> Callable[..., U]:
																def cache_decorator_wrapper(*UpperCAmelCase__     :		T   ) -> U:
																						if func not in cls.decorator_function_to_instance_map:
																												lowercase :							Dict    		=LRUCache(UpperCAmelCase__   )

																						lowercase :							Tuple    		=cls.decorator_function_to_instance_map[func].get(args[0]   )
																						if result is None:
																												lowercase :							Union[str, Any]    		=func(*UpperCAmelCase__   )
																												cls.decorator_function_to_instance_map[func].put(args[0]		,					UpperCAmelCase__   )
																						return result

																def cache_info() -> LRUCache[T, U]:
																						return cls.decorator_function_to_instance_map[func]

																setattr(UpperCAmelCase__		,					'''cache_info'''		,					UpperCAmelCase__   )  # noqa: B010

																return cache_decorator_wrapper

										return cache_decorator_inner


if __name__ == "__main__":
							import doctest

							doctest.testmod()








'''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			(   __magic_name__						:     int   , __magic_name__						:     int						)				->  Optional[int]:
						lowercase :							Optional[Any]    		=tuple((re.compile(x + '''$'''						) for x in qs)						)
						for i in range(len(__magic_name__						) - len(__magic_name__						) + 1						):
												lowercase :							Union[str, Any]    		=[x.match(__magic_name__						) for x, y in zip(__magic_name__   , ks[i:]						)]
												if matches and all(__magic_name__						):
																		return True
						return False






def 				_lowerCAmelCase			(   __magic_name__						:     Dict						)				->  List[str]:
						def replace(__magic_name__						:     Optional[int]   , __magic_name__						:     Optional[Any]						):
												for rule, replacement in rules:
																		if _match(__magic_name__   , __magic_name__						):
																								return replacement
												return val

						return replace






def 				_lowerCAmelCase			(   )				->  int:
						return [
						    # embeddings
						    (("transformer", "wpe", "embedding"), P('''mp'''   , __magic_name__						)),
						    (("transformer", "wte", "embedding"), P('''mp'''   , __magic_name__						)),
						    # atention
						    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(__magic_name__   , '''mp'''						)),
						    (("attention", "out_proj", "kernel"), P('''mp'''   , __magic_name__						)),
						    (("attention", "out_proj", "bias"), None),
						    # mlp
						    (("mlp", "c_fc", "kernel"), P(__magic_name__   , '''mp'''						)),
						    (("mlp", "c_fc", "bias"), P('''mp'''						)),
						    (("mlp", "c_proj", "kernel"), P('''mp'''   , __magic_name__						)),
						    (("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			(   __magic_name__						:     str						)				->  int:
						lowercase :							int    		=_get_partition_rules()
						lowercase :							Tuple    		=_replacement_rules(__magic_name__						)
						lowercase :							Any    		={k: _unmatched for k in flatten_dict(__magic_name__						)}
						lowercase :							Any    		={k: replace(__magic_name__   , __magic_name__						) for k, v in initd.items()}
						assert _unmatched not in result.values(), "Incomplete partition spec."
						return freeze(unflatten_dict(__magic_name__						)						)





