def keras_name_to_tf_name_stem_top(keras_name,
                                   use_ema=True,
                                   model_name_tf='efficientnet-b0'):
  """Mapping name in h5 to ckpt that is in stem or top (head).

  we map name keras_name that points to a weight in h5 file
  to a name of weight in ckpt file.

  Args:
    keras_name: str, the name of weight in the h5 file of keras implementation
    use_ema: Bool, use the ExponentialMovingAverage resuolt in ckpt or not
    model_name_tf: str, the name of model in ckpt.

  Returns:
    String for the name of weight as in ckpt file.

  Raises:
    KeyError: if we cannot parse the keras_name.
  """
  if use_ema:
    ema = '/ExponentialMovingAverage'
  else:
    ema = ''

  stem_top_dict = {
      'probs/bias:0': '{}/head/dense/bias{}',
      'probs/kernel:0': '{}/head/dense/kernel{}',
      'predictions/bias:0': '{}/head/dense/bias{}',
      'predictions/kernel:0': '{}/head/dense/kernel{}',
      'stem_conv/kernel:0': '{}/stem/conv2d/kernel{}',
      'top_conv/kernel:0': '{}/head/conv2d/kernel{}',
  }
  for x in stem_top_dict:
    stem_top_dict[x] = stem_top_dict[x].format(model_name_tf, ema)

  # stem batch normalization
  for bn_weights in ['beta', 'gamma', 'moving_mean', 'moving_variance']:
    tf_name = '{}/stem/tpu_batch_normalization/{}{}'.format(
        model_name_tf, bn_weights, ema)
    stem_top_dict['stem_bn/{}:0'.format(bn_weights)] = tf_name

  # top / head batch normalization
  for bn_weights in ['beta', 'gamma', 'moving_mean', 'moving_variance']:
    tf_name = '{}/head/tpu_batch_normalization/{}{}'.format(
        model_name_tf, bn_weights, ema)
    stem_top_dict['top_bn/{}:0'.format(bn_weights)] = tf_name

  if keras_name in stem_top_dict:
    return stem_top_dict[keras_name]
  raise KeyError('{} from h5 file cannot be parsed'.format(keras_name))