





import argparse


snake_case								=							"""docs/source/_static/js/custom.js"""

def 				SCREAMING_SNAKE_CASE__							(	snake_case__	:List[str]  )   ->    Union[str, Any]:
  with open(snake_case__      ,				encoding='utf-8'      ,				newline='\n'  ) as f:
    _lowercase						    =			f.readlines()
  _lowercase						    =			0

  # First let's put the right version
  while not lines[index].startswith('const stableVersion ='  ):
    index += 1
  _lowercase						    =			F"""const stableVersion = \"v{version}\"\n"""

  # Then update the dictionary
  while not lines[index].startswith('const versionMapping = {'  ):
    index += 1

  # We go until the end
  while not lines[index].startswith('}'  ):
    index += 1
  # We add the new version at the end
  lines[index - 1] += F"""    \"v{version}\": \"v{version}\",\n"""

  with open(snake_case__      ,				'w'      ,				encoding='utf-8'      ,				newline='\n'  ) as f:
    f.writelines(snake_case__  )


if __name__ == "__main__":
      snake_case								=							argparse.ArgumentParser()
      parser.add_argument("""--version""", help="""Release version.""")
      snake_case								=							parser.parse_args()
      update_custom_js(args.version)






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
_A					:				Optional[int]				=			object()

# For specifying empty leaf dict `{}`
_A					:				Tuple				=			object()


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


def 							__snake_case							(    lowerCAmelCase_	)				->						int:
    def replace(lowerCAmelCase_ ,							lowerCAmelCase_	):
        for rule, replacement in rules:
            if _match(lowerCAmelCase_ ,							lowerCAmelCase_	):
                return replacement
        return val

    return replace


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