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

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



def 						_lowerCAmelCase				(				UpperCamelCase_  ,				UpperCamelCase_				):
		__SCREAMING_SNAKE_CASE            = tuple((re.compile(x + """$"""				) for x in qs)				)
		for i in range(len(UpperCamelCase_				) - len(UpperCamelCase_				) + 1				):
				__SCREAMING_SNAKE_CASE            = [x.match(UpperCamelCase_				) for x, y in zip(UpperCamelCase_  ,				ks[i:]				)]
				if matches and all(UpperCamelCase_				):
						return True
		return False



def 						_lowerCAmelCase				(				UpperCamelCase_				):
		def replace(UpperCamelCase_  ,				UpperCamelCase_				):
				for rule, replacement in rules:
						if _match(UpperCamelCase_  ,				UpperCamelCase_				):
								return replacement
				return val

		return replace



def 						_lowerCAmelCase				(				):
		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_				):
		__SCREAMING_SNAKE_CASE            = _get_partition_rules()
		__SCREAMING_SNAKE_CASE            = _replacement_rules(UpperCamelCase_				)
		__SCREAMING_SNAKE_CASE            = {k: _unmatched for k in flatten_dict(UpperCamelCase_				)}
		__SCREAMING_SNAKE_CASE            = {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 						_lowerCAmelCase				(				UpperCamelCase_				):

		if isinstance(UpperCamelCase_  ,				UpperCamelCase_				):
				raise TypeError("""'float' object cannot be interpreted as an integer"""				)
		if isinstance(UpperCamelCase_  ,				UpperCamelCase_				):
				raise TypeError("""'str' object cannot be interpreted as an integer"""				)

		if num == 0:
				return "0b0"

		__SCREAMING_SNAKE_CASE            = False

		if num < 0:
				__SCREAMING_SNAKE_CASE            = True
				__SCREAMING_SNAKE_CASE            = -num

		__SCREAMING_SNAKE_CASE            = []
		while num > 0:
				binary.insert(0  ,				num % 2				)
				num >>= 1

		if negative:
				return "-0b" + "".join(str(UpperCamelCase_				) for e in binary				)

		return "0b" + "".join(str(UpperCamelCase_				) for e in binary				)


if __name__ == "__main__":
							import doctest

							doctest.testmod()




