Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/framework/kernels.py: 38%

1# Copyright 2018 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"""Functions for querying registered kernels.""" 

16 

17from tensorflow.core.framework import kernel_def_pb2 

18from tensorflow.python.client import pywrap_tf_session as c_api 

19from tensorflow.python.util import compat 

20 

21 

22def get_all_registered_kernels(): 

23 """Returns a KernelList proto of all registered kernels. 

24 """ 

25 buf = c_api.TF_GetAllRegisteredKernels() 

26 data = c_api.TF_GetBuffer(buf) 

27 kernel_list = kernel_def_pb2.KernelList() 

28 kernel_list.ParseFromString(compat.as_bytes(data)) 

29 return kernel_list 

30 

31 

32def get_registered_kernels_for_op(name): 

33 """Returns a KernelList proto of registered kernels for a given op. 

34 

35 Args: 

36 name: A string representing the name of the op whose kernels to retrieve. 

37 """ 

38 buf = c_api.TF_GetRegisteredKernelsForOp(name) 

39 data = c_api.TF_GetBuffer(buf) 

40 kernel_list = kernel_def_pb2.KernelList() 

41 kernel_list.ParseFromString(compat.as_bytes(data)) 

42 return kernel_list