def _get_output_to_inputs_map(self, output_signature):
# TODO(b/160550490): Remove local import.
from tensorflow_transform import graph_tools # pylint: disable=g-import-not-at-top
result = {}
for name, output in six.iteritems(output_signature):
components = self._get_component_tensors(output)
sinks = [self._as_operation(component) for component in components]
result[name] = graph_tools.retrieve_sources(sinks)
return result
def _get_output_to_inputs_map(self, output_signature):
"""Get all graph inputs that the tensors in output_signature depend on."""
result = {}
for name, output in six.iteritems(output_signature):
components = self._get_component_tensors(output)
sinks = [self._as_operation(component) for component in components]
# Ignore control dependencies when walking the graph as we only care about
# which user defined inputs this output depends on.
result[name] = graph_tools.retrieve_sources(
sinks, ignore_control_dependencies=True)
return result
def _get_output_to_inputs_map(
output_signature: Mapping[str, common_types.TensorType]
) -> Dict[str, Iterable[tf.Tensor]]:
"""Gets all graph inputs that the tensors in output_signature depend on."""
result = {}
for name, output in output_signature.items():
components = _get_component_tensors(output)
sinks = [_as_operation(component) for component in components]
# Ignore control dependencies when walking the graph as we only care about
# which user defined inputs this output depends on.
result[name] = graph_tools.retrieve_sources(
sinks, ignore_control_dependencies=True)
return result
def _get_output_to_inputs_map(self, output_signature):
"""Get all graph inputs that the tensors in output_signature depend on."""
# TODO(b/160550490): Remove local import.
from tensorflow_transform import graph_tools # pylint: disable=g-import-not-at-top
result = {}
for name, output in six.iteritems(output_signature):
components = self._get_component_tensors(output)
sinks = [self._as_operation(component) for component in components]
# Ignore control dependencies when walking the graph as we only care about
# which user defined inputs this output depends on.
result[name] = graph_tools.retrieve_sources(
sinks, ignore_control_dependencies=True)
return result
