Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/ops/handle_data_util.py: 34%
32 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
1# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Decorator to overrides the gradient for a function."""
17from tensorflow.python.client import pywrap_tf_session
18from tensorflow.python.framework import cpp_shape_inference_pb2
19from tensorflow.python.framework import dtypes
20from tensorflow.python.types import core
21from tensorflow.python.util import compat
24def get_resource_handle_data(graph_op):
25 assert (isinstance(graph_op, core.Symbol)
26 and not isinstance(graph_op, core.Value))
28 with graph_op.graph._c_graph.get() as c_graph: # pylint: disable=protected-access
29 handle_data = pywrap_tf_session.GetHandleShapeAndType(
30 c_graph, graph_op._as_tf_output()) # pylint: disable=protected-access
32 return cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData.FromString(
33 compat.as_bytes(handle_data))
36def get_handle_data(source_t):
37 """Obtains HandleData from a tensor."""
38 if isinstance(source_t, core.Value):
39 return source_t._handle_data # pylint: disable=protected-access
40 return get_resource_handle_data(source_t)
43def copy_handle_data(source_t, target_t):
44 """Copies HandleData for variant and resource type tensors if available.
46 The CppShapeInferenceResult::HandleData proto contains information about the
47 shapes and types of the element tensors of resource/variant type tensors.
48 We need to copy this across function boundaries, i.e., when capturing a
49 placeholder or when returning a function tensor as output. If we don't do this
50 the element tensors will have unknown shapes, e.g., if a TensorList variant
51 tensor is captured as a placeholder, elements popped from that list would have
52 unknown shape.
54 Args:
55 source_t: The tensor to copy HandleData from.
56 target_t: The tensor to copy HandleData to.
57 """
58 if (target_t.dtype == dtypes.resource or
59 target_t.dtype == dtypes.variant):
60 handle_data = get_handle_data(source_t)
61 set_handle_data(target_t, handle_data)
64def set_handle_data(target_t, handle_data):
65 """Sets handle data on the giver tensor."""
66 if (
67 handle_data is None
68 or not handle_data.is_set
69 or not handle_data.shape_and_type
70 ):
71 return
73 # pylint: disable=protected-access
74 if isinstance(target_t, core.Value):
75 target_t._handle_data = handle_data
76 return
77 with target_t.graph._c_graph.get() as c_graph:
78 pywrap_tf_session.SetHandleShapeAndType(c_graph, target_t._as_tf_output(),
79 handle_data.SerializeToString())
80 # pylint: enable=protected-access
83def create_handle_data(shape, dtype):
84 handle_data = cpp_shape_inference_pb2.CppShapeInferenceResult.HandleData()
85 handle_data.is_set = True
86 handle_data.shape_and_type.append(
87 cpp_shape_inference_pb2.CppShapeInferenceResult.HandleShapeAndType(
88 shape=shape.as_proto(), dtype=dtype.as_datatype_enum))
89 return handle_data