def supply_missing_inputs(
    structured_inputs: Mapping[str, common_types.TensorType],
    batch_size: int,
    missing_keys: Optional[Collection[str]] = None
) -> Mapping[str, common_types.TensorType]:
  """Supply inputs for unfed features.

  Supports only tf.Tensor, tf.SparseTensor and tf.RaggedTensor.

  Note: Since this returns placeholders, it should be called from within a graph
  context.

  Args:
    structured_inputs: a dict from keys to batches of placeholder graph tensors.
    batch_size: an integer representing the size of the batch returned.
    missing_keys: (Optional) a subset of the keys of `structured_inputs` for
      which concrete tensors need to be supplied. If `None`, tensors are
      supplied for all keys.

  Returns:
    A batch of placeholders with default values having the same structure as in
    `structured_inputs` for the keys in `missing_keys`.
  """
  missing_keys = missing_keys or list(structured_inputs)
  # Return placeholders to ensure that tensor shape is not constrained to the
  # dummy shape of the missing tensor created here during tracing.
  result = {}
  for key in missing_keys:
    tensor = structured_inputs[key]
    if isinstance(tensor, tf.Tensor) or (isinstance(tensor, tf.RaggedTensor) and
                                         tensor.ragged_rank == 0):
      missing_tensor = supply_missing_tensor(batch_size, tensor.shape,
                                             tensor.dtype)
      result[key] = tf.raw_ops.PlaceholderWithDefault(
          input=missing_tensor, shape=tensor.shape)
    elif isinstance(tensor, tf.SparseTensor):
      values = supply_missing_tensor(batch_size, tensor.values.shape,
                                     tensor.values.dtype)
      dense_rank = tensor.shape.ndims
      # Since values is always a 1-D tensor, set index for every ith value in
      # values to be [i 0 0 ...]. Each index should be compatible with the
      # rank of the SparseTensor. Hence, the number of 0s is dense_rank-1.
      actual_batch_size = tf.shape(values)[0]
      indices = tf.stack(
          [tf.range(actual_batch_size, dtype=tf.int64)] +
          [tf.zeros(actual_batch_size, dtype=tf.int64)] * (dense_rank - 1),
          axis=1)
      dense_shape = tf.cast(
          [actual_batch_size] + [1] * (dense_rank - 1), dtype=tf.int64)

      indices = tf.raw_ops.PlaceholderWithDefault(
          input=indices, shape=tensor.indices.shape)
      values = tf.raw_ops.PlaceholderWithDefault(
          input=values, shape=tensor.values.shape)
      dense_shape = tf.raw_ops.PlaceholderWithDefault(
          input=dense_shape, shape=tensor.dense_shape.shape)
      result[key] = tf.SparseTensor(
          indices=indices, values=values, dense_shape=dense_shape)
    elif isinstance(tensor, tf.RaggedTensor):
      # Builds a ragged tensor similar to tf.ragged.placeholder, except with
      # default values for all components.
      ragged_rank = tensor.ragged_rank
      values = supply_missing_tensor(batch_size, tensor.flat_values.shape,
                                     tensor.flat_values.dtype)
      result[key] = tf.raw_ops.PlaceholderWithDefault(
          input=values, shape=tensor.flat_values.shape)
      for _ in range(ragged_rank):
        if isinstance(values, tf.RaggedTensor):
          values_batch_size = values.bounding_shape(axis=0)
        else:
          values_batch_size = tf.shape(values)[0]
        row_splits = tf.range(values_batch_size + 1, dtype=tf.int64)
        values = tf.RaggedTensor.from_row_splits(
            values, row_splits, validate=False)
        row_splits = tf.raw_ops.PlaceholderWithDefault(
            input=row_splits, shape=[None])
        result[key] = tf.RaggedTensor.from_row_splits(
            result[key], row_splits, validate=False)
    else:
      raise ValueError('Received unsupported input tensor type. Only '
                       'dense/sparse/ragged tensors are currently supported.')
  return result
Supply inputs for unfed features.

  Supports only tf.Tensor, tf.SparseTensor and tf.RaggedTensor.

  Note: Since this returns placeholders, it should be called from within a graph
  context.

  Args:
    structured_inputs: a dict from keys to batches of placeholder graph tensors.
    batch_size: an integer representing the size of the batch returned.
    missing_keys: (Optional) a subset of the keys of `structured_inputs` for
      which concrete tensors need to be supplied. If `None`, tensors are
      supplied for all keys.

  Returns:
    A batch of placeholders with default values having the same structure as in
    `structured_inputs` for the keys in `missing_keys`.
  
tensorflow_transform/tf2_utils.py