def _apply_v2_transform_model(
      self, logical_input_map: Mapping[str, common_types.TensorType]
  ) -> Dict[str, common_types.TensorType]:
    """Applies a V2 transform graph to dictionary of (Composite)Tensors.

    This method applies the transformation graph to the `logical_input_map` to
    return only outputs that can be computed from the keys provided in
    `logical_input_map`.

    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.

    Returns:
      A dict of logical name to Tensor, as provided by the output signature of
      the transform graph.
    """

    unfed_input_keys = self._get_unfed_input_keys(logical_input_map.keys())
    feeds = self._get_feeds(unfed_input_keys)
    modified_inputs = self._format_input_map_as_tensors(logical_input_map)

    if unfed_input_keys:
      batch_size = 1
      if logical_input_map:
        an_input = next(iter(logical_input_map.values()))
        if isinstance(an_input, tf.RaggedTensor):
          batch_size = an_input.bounding_shape(axis=0)
        elif tf.shape(an_input)[0] is not None:
          batch_size = tf.shape(an_input)[0]

      missing_inputs = self._get_missing_inputs(unfed_input_keys, batch_size)
      modified_inputs.update(missing_inputs)

    flattened_inputs = tf.nest.flatten(modified_inputs, expand_composites=True)

    # self._wrapped_function.inputs may be longer than flattened_inputs as it
    # also contains captured inputs. However, we only want the user inputs here
    # so we don't assert equal length.
    for input_t, wrapped_input in zip(flattened_inputs,
                                      self._wrapped_function.inputs):
      try:
        wrapped_input.shape.assert_is_compatible_with(input_t.shape)
      except ValueError as e:
        raise ValueError('{}: {}'.format(input_t, e))

    transformed_features = self._wrapped_function(*flattened_inputs)
    fetches_keys = self._get_fetches_keys(feeds)
    return {key: transformed_features[key] for key in fetches_keys}
Applies a V2 transform graph to dictionary of (Composite)Tensors.

    This method applies the transformation graph to the `logical_input_map` to
    return only outputs that can be computed from the keys provided in
    `logical_input_map`.

    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.

    Returns:
      A dict of logical name to Tensor, as provided by the output signature of
      the transform graph.
    
tensorflow_transform/saved/saved_transform_io_v2.py