Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/ops/resources.py: 47%

1# Copyright 2016 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 

16 

17"""Utilities for using generic resources.""" 

18# pylint: disable=g-bad-name 

19import collections 

20import os 

21 

22from tensorflow.python.framework import dtypes 

23from tensorflow.python.framework import ops 

24from tensorflow.python.ops import array_ops 

25from tensorflow.python.ops import array_ops_stack 

26from tensorflow.python.ops import control_flow_ops 

27from tensorflow.python.ops import math_ops 

28from tensorflow.python.util import tf_should_use 

29 

30 

31_Resource = collections.namedtuple("_Resource", 

32 ["handle", "create", "is_initialized"]) 

33 

34 

35def register_resource(handle, create_op, is_initialized_op, is_shared=True): 

36 """Registers a resource into the appropriate collections. 

37 

38 This makes the resource findable in either the shared or local resources 

39 collection. 

40 

41 Args: 

42 handle: op which returns a handle for the resource. 

43 create_op: op which initializes the resource. 

44 is_initialized_op: op which returns a scalar boolean tensor of whether 

45 the resource has been initialized. 

46 is_shared: if True, the resource gets added to the shared resource 

47 collection; otherwise it gets added to the local resource collection. 

48 

49 """ 

50 resource = _Resource(handle, create_op, is_initialized_op) 

51 if is_shared: 

52 ops.add_to_collection(ops.GraphKeys.RESOURCES, resource) 

53 else: 

54 ops.add_to_collection(ops.GraphKeys.LOCAL_RESOURCES, resource) 

55 

56 

57def shared_resources(): 

58 """Returns resources visible to all tasks in the cluster.""" 

59 return ops.get_collection(ops.GraphKeys.RESOURCES) 

60 

61 

62def local_resources(): 

63 """Returns resources intended to be local to this session.""" 

64 return ops.get_collection(ops.GraphKeys.LOCAL_RESOURCES) 

65 

66 

67def report_uninitialized_resources(resource_list=None, 

68 name="report_uninitialized_resources"): 

69 """Returns the names of all uninitialized resources in resource_list. 

70 

71 If the returned tensor is empty then all resources have been initialized. 

72 

73 Args: 

74 resource_list: resources to check. If None, will use shared_resources() + 

75 local_resources(). 

76 name: name for the resource-checking op. 

77 

78 Returns: 

79 Tensor containing names of the handles of all resources which have not 

80 yet been initialized. 

81 

82 """ 

83 if resource_list is None: 

84 resource_list = shared_resources() + local_resources() 

85 with ops.name_scope(name): 

86 # Run all operations on CPU 

87 local_device = os.environ.get( 

88 "TF_DEVICE_FOR_UNINITIALIZED_VARIABLE_REPORTING", "/cpu:0") 

89 with ops.device(local_device): 

90 if not resource_list: 

91 # Return an empty tensor so we only need to check for returned tensor 

92 # size being 0 as an indication of model ready. 

93 return array_ops.constant([], dtype=dtypes.string) 

94 # Get a 1-D boolean tensor listing whether each resource is initialized. 

95 variables_mask = math_ops.logical_not( 

96 array_ops_stack.stack([r.is_initialized for r in resource_list])) 

97 # Get a 1-D string tensor containing all the resource names. 

98 variable_names_tensor = array_ops.constant( 

99 [s.handle.name for s in resource_list]) 

100 # Return a 1-D tensor containing all the names of uninitialized resources. 

101 return array_ops.boolean_mask(variable_names_tensor, variables_mask) 

102 

103 

104@tf_should_use.should_use_result 

105def initialize_resources(resource_list, name="init"): 

106 """Initializes the resources in the given list. 

107 

108 Args: 

109 resource_list: list of resources to initialize. 

110 name: name of the initialization op. 

111 

112 Returns: 

113 op responsible for initializing all resources. 

114 """ 

115 if resource_list: 

116 return control_flow_ops.group(*[r.create for r in resource_list], name=name) 

117 return control_flow_ops.no_op(name=name)