






"""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
__lowerCamelCase							=  object()

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



def   a    (     __UpperCAmelCase      :    int      ,   __UpperCAmelCase      :    Tuple  )    ->							int:
						__magic_name__:	str							=       tuple((re.compile(x + """$"""  ) for x in qs)  )
						for i in range(len(__UpperCAmelCase  ) - len(__UpperCAmelCase  ) + 1  ):
												__magic_name__:	Dict							=       [x.match(__UpperCAmelCase  ) for x, y in zip(__UpperCAmelCase      ,   ks[i:]  )]
												if matches and all(__UpperCAmelCase  ):
																		return True
						return False



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

						return replace



def   a    (     )    ->							str:
						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   a    (     __UpperCAmelCase      :    Union[str, Any]  )    ->							Dict:
						__magic_name__:	Optional[Any]							=       _get_partition_rules()
						__magic_name__:	Any							=       _replacement_rules(__UpperCAmelCase  )
						__magic_name__:	int							=       {k: _unmatched for k in flatten_dict(__UpperCAmelCase  )}
						__magic_name__:	int							=       {k: replace(__UpperCAmelCase      ,   __UpperCAmelCase  ) for k, v in initd.items()}
						assert _unmatched not in result.values(), "Incomplete partition spec."
						return freeze(unflatten_dict(__UpperCAmelCase  )  )











"""simple docstring"""





def   a    (     __UpperCAmelCase      :    int = 1_0_0  )    ->							int:
						__magic_name__:	str							=       0
						__magic_name__:	Any							=       0
						for i in range(1      ,   n + 1  ):
												sum_of_squares += i**2
												sum_of_ints += i
						return sum_of_ints**2 - sum_of_squares


if __name__ == "__main__":
							print(f'''{solution() = }''')



