def convert_tf_config_to_jax_bert(config):
  """Convert TF BERT model config to be compatible with JAX BERT model.

  Args:
    config: dictionary of TF model configurations

  Returns:
    dictionary of param names and values compatible with JAX BERT model
  """
  unnecessary_keys = ['initializer_range', 'backward_compatible',
                      'embedding_size']
  for key in unnecessary_keys:
    if key in config:
      config.pop(key)

  # change TF parameter names to match JAX parameter names
  mapping = {
      'attention_dropout_rate': 'attention_probs_dropout_prob',
      'hidden_activation': 'hidden_act',
      'dropout_rate': 'hidden_dropout_prob',
      'emb_dim': 'hidden_size',
      'mlp_dim': 'intermediate_size',
      'max_len': 'max_position_embeddings',
      'num_heads': 'num_attention_heads',
      'num_layers': 'num_hidden_layers'
  }
  for jax_key, tf_key in mapping.items():
    config[jax_key] = config.pop(tf_key)

  return config