def map_per_key_reductions(tensors_to_map: Tuple[tf.Tensor, ...],
                           key: common_types.TensorType, key_vocab: tf.Tensor,
                           original_input: common_types.TensorType,
                           reduce_instance_dims: bool) -> Tuple[tf.Tensor, ...]:
  """Rearrange the reduced per-key result to correspond to the original keys.

  Args:
    tensors_to_map: A tuple of 1-D `Tensor`s that are same shape as key_vocab,
      to be mapped to respective key.
    key: A `Tensor` or `CompositeTensor`.
    key_vocab: A 1-D `Tensor`.
    original_input: A `Tensor` or `CompositeTensor`.
    reduce_instance_dims: A `bool`. True if tensors_to_map are reduced in
      dimension, else False.

  Returns:
    A tuple same length as tensors_to_map, of `Tensor`s the same dimension as
    original_input. We are mapping using the key for each original_input,
    but output rank needs to match original_input in the dense case.
    For the sparse case, it is enough for output to match original_input.values.
    Any missing key would result in a mapping to 0.
  """

  _, key = _validate_and_get_dense_value_key_inputs(original_input, key)
  key_indices = lookup_key(key, key_vocab)

  ndims = (None if isinstance(original_input,
                              (tf.SparseTensor, tf.RaggedTensor)) else
           original_input.get_shape().ndims)

  # Append 0s to allow mapping OOVs to it.
  tensors_to_map = [
      tf.concat([t, tf.expand_dims(tf.zeros_like(t[0]), 0)], axis=0)
      for t in tensors_to_map
  ]

  # Replace `-1`s due to OOV with size of key_vocab.
  adjusted_indices = tf.where(
      key_indices >= 0, key_indices,
      tf.cast(
          tf.fill(tf.shape(key_indices), tf.size(key_vocab)), dtype=tf.int64))
  axis = -1 if reduce_instance_dims else 0
  mapped_result = [
      _align_dims(tf.gather(t, adjusted_indices, axis=axis), ndims)
      for t in tensors_to_map
  ]

  return tuple(mapped_result)
Rearrange the reduced per-key result to correspond to the original keys.

  Args:
    tensors_to_map: A tuple of 1-D `Tensor`s that are same shape as key_vocab,
      to be mapped to respective key.
    key: A `Tensor` or `CompositeTensor`.
    key_vocab: A 1-D `Tensor`.
    original_input: A `Tensor` or `CompositeTensor`.
    reduce_instance_dims: A `bool`. True if tensors_to_map are reduced in
      dimension, else False.

  Returns:
    A tuple same length as tensors_to_map, of `Tensor`s the same dimension as
    original_input. We are mapping using the key for each original_input,
    but output rank needs to match original_input in the dense case.
    For the sparse case, it is enough for output to match original_input.values.
    Any missing key would result in a mapping to 0.
  
tensorflow_transform/tf_utils.py