'''simple docstring'''




import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


__UpperCAmelCase					=logging.get_logger(__name__)






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       int:
    __lowerCamelCase            =    r'''\w+[.]\d+'''
    __lowerCamelCase            =    re.findall(UpperCamelCase__							,    UpperCamelCase__							)
    for pat in pats:
        __lowerCamelCase            =    key.replace(UpperCamelCase__							,    '''_'''.join(pat.split('''.'''							)							)							)
    return key






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)      ->       Dict:
    __lowerCamelCase            =    pt_tuple_key[:-1] + ('''scale''',)
    if (
        any('''norm''' in str_ for str_ in pt_tuple_key							)
        and (pt_tuple_key[-1] == "bias")
        and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
        and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
    ):
        __lowerCamelCase            =    pt_tuple_key[:-1] + ('''scale''',)
        return renamed_pt_tuple_key, pt_tensor
    elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
        __lowerCamelCase            =    pt_tuple_key[:-1] + ('''scale''',)
        return renamed_pt_tuple_key, pt_tensor

    # embedding
    if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
        __lowerCamelCase            =    pt_tuple_key[:-1] + ('''embedding''',)
        return renamed_pt_tuple_key, pt_tensor

    # conv layer
    __lowerCamelCase            =    pt_tuple_key[:-1] + ('''kernel''',)
    if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
        __lowerCamelCase            =    pt_tensor.transpose(2							,    3							,    1							,    0							)
        return renamed_pt_tuple_key, pt_tensor

    # linear layer
    __lowerCamelCase            =    pt_tuple_key[:-1] + ('''kernel''',)
    if pt_tuple_key[-1] == "weight":
        __lowerCamelCase            =    pt_tensor.T
        return renamed_pt_tuple_key, pt_tensor

    # old PyTorch layer norm weight
    __lowerCamelCase            =    pt_tuple_key[:-1] + ('''weight''',)
    if pt_tuple_key[-1] == "gamma":
        return renamed_pt_tuple_key, pt_tensor

    # old PyTorch layer norm bias
    __lowerCamelCase            =    pt_tuple_key[:-1] + ('''bias''',)
    if pt_tuple_key[-1] == "beta":
        return renamed_pt_tuple_key, pt_tensor

    return pt_tuple_key, pt_tensor






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__=42							)      ->       int:
    # Step 1: Convert pytorch tensor to numpy
    __lowerCamelCase            =    {k: v.numpy() for k, v in pt_state_dict.items()}

    # Step 2: Since the model is stateless, get random Flax params
    __lowerCamelCase            =    flax_model.init_weights(PRNGKey(UpperCamelCase__							)							)

    __lowerCamelCase            =    flatten_dict(UpperCamelCase__							)
    __lowerCamelCase            =    {}

    # Need to change some parameters name to match Flax names
    for pt_key, pt_tensor in pt_state_dict.items():
        __lowerCamelCase            =    rename_key(UpperCamelCase__							)
        __lowerCamelCase            =    tuple(renamed_pt_key.split('''.'''							)							)

        # Correctly rename weight parameters
        __lowerCamelCase      ,    __lowerCamelCase            =    rename_key_and_reshape_tensor(UpperCamelCase__							,    UpperCamelCase__							,    UpperCamelCase__							)

        if flax_key in random_flax_state_dict:
            if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
                raise ValueError(
                    f"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
                    f"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""							)

        # also add unexpected weight so that warning is thrown
        __lowerCamelCase            =    jnp.asarray(UpperCamelCase__							)

    return unflatten_dict(UpperCamelCase__							)

'''simple docstring'''




def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							)      ->       float:
    return base * power(UpperCamelCase__							,    (exponent - 1)							) if exponent else 1


if __name__ == "__main__":
  print("Raise base to the power of exponent using recursion...")
  __UpperCAmelCase					=int(input("Enter the base: ").strip())
  __UpperCAmelCase					=int(input("Enter the exponent: ").strip())
  __UpperCAmelCase					=power(base, abs(exponent))
  if exponent < 0:  # power() does not properly deal w/ negative exponents
    __UpperCAmelCase					=1 / result
  print(f'{base} to the power of {exponent} is {result}')
