def _expand_input_map(logical_input_map, input_signature):
  """Expands user provided inputs to component tensors in the graph.

  The user specified `logical_input_map` contains mappings from logical feature
  names to `Tensor`s or `CompositeTensor`s. These are expanded into mappings
  from component tensor names in the graph to their corresponding component
  tensor value.

  Args:
    logical_input_map: a dict of logical name to Tensor.  The logical names must
      be a subset of those in the input signature of the transform graph, and
      the corresponding Tensors must have the expected types and shapes.
    input_signature: The inputs of a `SignatureDef` proto for the graph to be
      imported.

  Returns:
    A map from tensor names in `input_signature` to the tensors
    specified in `logical_input_map`.
  """
  result = {}
  for logical_name, replacement in logical_input_map.items():
    tensor_info = input_signature[logical_name]
    encoding = tensor_info.WhichOneof('encoding')
    if encoding == 'coo_sparse':
      assert isinstance(replacement, tf.SparseTensor), logical_name
      result[tensor_info.coo_sparse.indices_tensor_name] = replacement.indices
      result[tensor_info.coo_sparse.values_tensor_name] = replacement.values
      result[tensor_info.coo_sparse.dense_shape_tensor_name] = (
          replacement.dense_shape)
    elif encoding == 'composite_tensor':
      component_infos = tensor_info.composite_tensor.components
      component_tensors = tf.nest.flatten(replacement, expand_composites=True)
      for (info, tensor) in zip(component_infos, component_tensors):
        result[info.name] = tensor
    elif encoding == 'name':
      result[tensor_info.name] = replacement
    else:
      raise ValueError('Unsupported TensorInfo encoding %s' % encoding)
  return result
Expands user provided inputs to component tensors in the graph.

  The user specified `logical_input_map` contains mappings from logical feature
  names to `Tensor`s or `CompositeTensor`s. These are expanded into mappings
  from component tensor names in the graph to their corresponding component
  tensor value.

  Args:
    logical_input_map: a dict of logical name to Tensor.  The logical names must
      be a subset of those in the input signature of the transform graph, and
      the corresponding Tensors must have the expected types and shapes.
    input_signature: The inputs of a `SignatureDef` proto for the graph to be
      imported.

  Returns:
    A map from tensor names in `input_signature` to the tensors
    specified in `logical_input_map`.
  
tensorflow_transform/saved/saved_transform_io.py