def _get_rename_keys(features):
  """Returns a function that converts TensorDicts to TensorDictTriples.

  The triple contains three TensorDicts: `(inputs, labels, weights)`. Labels
  and weights correspond to keys with `/value` and `/weight` suffixes,
  respectively. All other keys are considered input keys. The `image/encoded`
  key in `features` is renamed to `IMAGE_KEY`.
  Args:
    features: The features dictionary used by `tf.io.parse_single_example`.
  Returns:
    A function that converts a single TensorDict into separate inputs, labels,
    and weights TensorDicts.
  """
  weight_keys = {key for key in features if key.endswith('/weight')}
  label_keys = {key for key in features if key.endswith('/value')}
  input_keys = set(features.keys()) - weight_keys - label_keys

  def _rename_keys(example):
    inputs = {}
    labels = {}
    weights = {}

    for key, value in example.items():
      key_split = key.split('/')
      new_key = key_split[1]

      if key in weight_keys:
        weights[new_key] = value
      elif key in label_keys:
        labels[new_key] = value
      elif key in input_keys:
        new_key = 'image' if new_key == 'encoded' else new_key
        inputs[new_key] = value
      else:
        raise ValueError(f'Unexpected key: {key}.')
    return inputs, labels, weights

  return _rename_keys