1# Copyright 2015 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"""Global registry for OpDefs.""" 17 18import threading 19 20from tensorflow.core.framework import op_def_pb2 21# pylint: disable=invalid-import-order,g-bad-import-order, wildcard-import, unused-import 22from tensorflow.python import pywrap_tensorflow 23from tensorflow.python.framework import _op_def_registry 24 25# The cache amortizes ProtoBuf serialization/deserialization overhead 26# on the language boundary. If an OpDef has been looked up, its Python 27# representation is cached. 28_cache = {} 29_cache_lock = threading.Lock() 30 31 32def get(name): 33 """Returns an OpDef for a given `name` or None if the lookup fails.""" 34 try: 35 return _cache[name] 36 except KeyError: 37 pass 38 39 with _cache_lock: 40 try: 41 # Return if another thread has already populated the cache. 42 return _cache[name] 43 except KeyError: 44 pass 45 46 serialized_op_def = _op_def_registry.get(name) 47 if serialized_op_def is None: 48 return None 49 50 op_def = op_def_pb2.OpDef() 51 op_def.ParseFromString(serialized_op_def) 52 _cache[name] = op_def 53 return op_def 54 55 56# TODO(b/141354889): Remove once there are no callers. 57def sync(): 58 """No-op. Used to synchronize the contents of the Python registry with C++.""" 59