def _finalize_wrapped_function(
      self, unfed_input_keys: Iterable[str],
      fetches_keys: Iterable[str]) -> function.ConcreteFunction:
    """Constructs a function that can be invoked without `unfed_input_keys`."""
    original_input_signature = (
        self._wrapped_function.structured_input_signature[0][0])
    input_signature = {
        k: v
        for k, v in original_input_signature.items()
        if k not in unfed_input_keys
    }

    @tf.function(input_signature=[input_signature], autograph=False)
    def wrapped_finalized(inputs):
      missing_inputs = self._get_missing_inputs(unfed_input_keys, batch_size=1)
      # Directly modifying inputs is not allowed in a tf.function. Hence, we
      # make a deep copy here.
      inputs_copy = tf_utils.copy_tensors(inputs)
      inputs_copy.update(missing_inputs)
      flattened_inputs = tf.nest.flatten(inputs_copy, expand_composites=True)
      transformed_features = self._wrapped_function(*flattened_inputs)
      return {key: transformed_features[key] for key in fetches_keys}

    return wrapped_finalized.get_concrete_function()
Constructs a function that can be invoked without `unfed_input_keys`.
tensorflow_transform/saved/saved_transform_io_v2.py