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"""Function for loading TensorFlow plugins.""" 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21import hashlib 22import imp 23import sys 24import threading # pylint: disable=unused-import 25 26from tensorflow.core.framework import op_def_pb2 27from tensorflow.core.lib.core import error_codes_pb2 # pylint: disable=unused-import 28from tensorflow.python import pywrap_tensorflow as py_tf 29from tensorflow.python.framework import errors_impl 30from tensorflow.python.util import compat 31from tensorflow.python.util.tf_export import tf_export 32 33 34@tf_export('load_op_library') 35def load_op_library(library_filename): 36 """Loads a TensorFlow plugin, containing custom ops and kernels. 37 38 Pass "library_filename" to a platform-specific mechanism for dynamically 39 loading a library. The rules for determining the exact location of the 40 library are platform-specific and are not documented here. When the 41 library is loaded, ops and kernels registered in the library via the 42 `REGISTER_*` macros are made available in the TensorFlow process. Note 43 that ops with the same name as an existing op are rejected and not 44 registered with the process. 45 46 Args: 47 library_filename: Path to the plugin. 48 Relative or absolute filesystem path to a dynamic library file. 49 50 Returns: 51 A python module containing the Python wrappers for Ops defined in 52 the plugin. 53 54 Raises: 55 RuntimeError: when unable to load the library or get the python wrappers. 56 """ 57 with errors_impl.raise_exception_on_not_ok_status() as status: 58 lib_handle = py_tf.TF_LoadLibrary(library_filename, status) 59 60 op_list_str = py_tf.TF_GetOpList(lib_handle) 61 op_list = op_def_pb2.OpList() 62 op_list.ParseFromString(compat.as_bytes(op_list_str)) 63 wrappers = py_tf.GetPythonWrappers(op_list_str) 64 65 # Delete the library handle to release any memory held in C 66 # that are no longer needed. 67 py_tf.TF_DeleteLibraryHandle(lib_handle) 68 69 # Get a unique name for the module. 70 module_name = hashlib.md5(wrappers).hexdigest() 71 if module_name in sys.modules: 72 return sys.modules[module_name] 73 module = imp.new_module(module_name) 74 # pylint: disable=exec-used 75 exec(wrappers, module.__dict__) 76 # Stash away the library handle for making calls into the dynamic library. 77 module.LIB_HANDLE = lib_handle 78 # OpDefs of the list of ops defined in the library. 79 module.OP_LIST = op_list 80 sys.modules[module_name] = module 81 return module 82 83 84@tf_export('load_file_system_library') 85def load_file_system_library(library_filename): 86 """Loads a TensorFlow plugin, containing file system implementation. 87 88 Pass `library_filename` to a platform-specific mechanism for dynamically 89 loading a library. The rules for determining the exact location of the 90 library are platform-specific and are not documented here. 91 92 Args: 93 library_filename: Path to the plugin. 94 Relative or absolute filesystem path to a dynamic library file. 95 96 Returns: 97 None. 98 99 Raises: 100 RuntimeError: when unable to load the library. 101 """ 102 with errors_impl.raise_exception_on_not_ok_status() as status: 103 lib_handle = py_tf.TF_LoadLibrary(library_filename, status) 104