# Path: flaxformer/param_conversion_util.py
# def load_tf_params(checkpoint_path: str) -> Dict[str, Any]:
# def convert_tf_params(
#     tf_params: Dict[str, Any],
#     mapping: Mapping[Tuple[str, ...], str],
#     tf_prefix: str = '',
# ) -> Dict[str, Any]:
# def get_int_regex_matches(pattern: str, state: Mapping[str, Any]) -> List[int]:


from typing import Any, Dict, Tuple
from flax.core import frozen_dict
from flaxformer import param_conversion_util
import flax
import tensorflow as tf

  Args:
    checkpoint_path: Path to the TF checkpoint.

  Returns:
    A tuple of frozen dicts that are the parameters for the bert.BertEncoder and
    bert.heads.BertPooler.
  """
  ckpt_reader = tf.train.load_checkpoint(checkpoint_path)
  tf_params = {
      tf_name: ckpt_reader.get_tensor(tf_name)
      for tf_name in ckpt_reader.get_variable_to_dtype_map()
  }
  encoder_params = convert_bert_encoder_params(tf_params, 'bert/')
  pooler_params = param_conversion_util.convert_tf_params(
      tf_params, {
          ('dense', 'bias'): 'bias',
          ('dense', 'kernel'): 'kernel'
      }, 'bert/pooler/dense/')

  # TODO: Update this when BertEncoder is deleted.
  return frozen_dict.freeze(encoder_params), frozen_dict.freeze(pooler_params)


# TODO: Delete this when BertEncoder is deleted.
def convert_bert_encoder_params(tf_params: Dict[str, tf.Tensor],
                                prefix: str) -> Dict[str, Any]:
  """Loads all Flax BertEncoder parameters from the TF params."""
  return {
      'embedder':

# Path: flaxformer/param_conversion_util.py
# def load_tf_params(checkpoint_path: str) -> Dict[str, Any]:
# def convert_tf_params(
#     tf_params: Dict[str, Any],
#     mapping: Mapping[Tuple[str, ...], str],
#     tf_prefix: str = '',
# ) -> Dict[str, Any]:
# def get_int_regex_matches(pattern: str, state: Mapping[str, Any]) -> List[int]:


# Path: flaxformer/architectures/bert/bert_checkpoint_converter.py
from typing import Any, Dict, Tuple
from flax.core import frozen_dict
from flaxformer import param_conversion_util
import flax
import tensorflow as tf


  Args:
    checkpoint_path: Path to the TF checkpoint.

  Returns:
    A tuple of frozen dicts that are the parameters for the bert.BertEncoder and
    bert.heads.BertPooler.
  """
  ckpt_reader = tf.train.load_checkpoint(checkpoint_path)
  tf_params = {
      tf_name: ckpt_reader.get_tensor(tf_name)
      for tf_name in ckpt_reader.get_variable_to_dtype_map()
  }
  encoder_params = convert_bert_encoder_params(tf_params, 'bert/')
  pooler_params = param_conversion_util.convert_tf_params(
      tf_params, {
          ('dense', 'bias'): 'bias',
          ('dense', 'kernel'): 'kernel'
      }, 'bert/pooler/dense/')

  # TODO: Update this when BertEncoder is deleted.
  return frozen_dict.freeze(encoder_params), frozen_dict.freeze(pooler_params)


# TODO: Delete this when BertEncoder is deleted.
def convert_bert_encoder_params(tf_params: Dict[str, tf.Tensor],
                                prefix: str) -> Dict[str, Any]:
  """Loads all Flax BertEncoder parameters from the TF params."""
  return {
      'embedder':

          param_conversion_util.convert_tf_params(