wehub-resource-sync
82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Helpers to traverse the Dataset dependency structure."""
|
|
import queue
|
|
|
|
from tensorflow.python.framework import dtypes
|
|
|
|
|
|
OP_TYPES_ALLOWLIST = ["DummyIterationCounter"]
|
|
# We allowlist all ops that produce variant tensors as output. This is a bit
|
|
# of overkill but the other dataset _inputs() traversal strategies can't
|
|
# cover the case of function inputs that capture dataset variants.
|
|
TENSOR_TYPES_ALLOWLIST = [dtypes.variant]
|
|
|
|
|
|
def _traverse(dataset, op_filter_fn):
|
|
"""Traverse a dataset graph, returning nodes matching `op_filter_fn`."""
|
|
result = []
|
|
bfs_q = queue.Queue()
|
|
bfs_q.put(dataset._variant_tensor.op) # pylint: disable=protected-access
|
|
visited = []
|
|
while not bfs_q.empty():
|
|
op = bfs_q.get()
|
|
visited.append(op)
|
|
if op_filter_fn(op):
|
|
result.append(op)
|
|
for i in op.inputs:
|
|
input_op = i.op
|
|
if input_op not in visited:
|
|
bfs_q.put(input_op)
|
|
return result
|
|
|
|
|
|
def obtain_capture_by_value_ops(dataset):
|
|
"""Given an input dataset, finds all allowlisted ops used for construction.
|
|
|
|
Allowlisted ops are stateful ops which are known to be safe to capture by
|
|
value.
|
|
|
|
Args:
|
|
dataset: Dataset to find allowlisted stateful ops for.
|
|
|
|
Returns:
|
|
A list of variant_tensor producing dataset ops used to construct this
|
|
dataset.
|
|
"""
|
|
|
|
def capture_by_value(op):
|
|
return (op.outputs[0].dtype in TENSOR_TYPES_ALLOWLIST or
|
|
op.type in OP_TYPES_ALLOWLIST)
|
|
|
|
return _traverse(dataset, capture_by_value)
|
|
|
|
|
|
def obtain_all_variant_tensor_ops(dataset):
|
|
"""Given an input dataset, finds all dataset ops used for construction.
|
|
|
|
A series of transformations would have created this dataset with each
|
|
transformation including zero or more Dataset ops, each producing a dataset
|
|
variant tensor. This method outputs all of them.
|
|
|
|
Args:
|
|
dataset: Dataset to find variant tensors for.
|
|
|
|
Returns:
|
|
A list of variant_tensor producing dataset ops used to construct this
|
|
dataset.
|
|
"""
|
|
return _traverse(dataset, lambda op: op.outputs[0].dtype == dtypes.variant)
|