1# Copyright 2020 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"""DLPack modules for Tensorflow.""" 16 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21from tensorflow.python import pywrap_tfe 22from tensorflow.python.eager import context 23from tensorflow.python.util.tf_export import tf_export 24 25 26@tf_export("experimental.dlpack.to_dlpack", v1=[]) 27def to_dlpack(tf_tensor): 28 """Returns the dlpack capsule representing the tensor. 29 30 This operation ensures the underlying data memory is ready when returns. 31 32 ```python 33 a = tf.tensor([1, 10]) 34 dlcapsule = tf.experimental.dlpack.to_dlpack(a) 35 # dlcapsule represents the dlpack data structure 36 ``` 37 38 Args: 39 tf_tensor: Tensorflow eager tensor, to be converted to dlpack capsule. 40 41 Returns: 42 A PyCapsule named as dltensor, which shares the underlying memory to other 43 framework. This PyCapsule can be consumed only once. 44 """ 45 return pywrap_tfe.TFE_ToDlpackCapsule(tf_tensor) 46 47 48@tf_export("experimental.dlpack.from_dlpack", v1=[]) 49def from_dlpack(dlcapsule): 50 """Returns the Tensorflow eager tensor. 51 52 The returned tensor uses the memory shared by dlpack capsules from other 53 framework. 54 55 ```python 56 a = tf.experimental.dlpack.from_dlpack(dlcapsule) 57 # `a` uses the memory shared by dlpack 58 ``` 59 60 Args: 61 dlcapsule: A PyCapsule named as dltensor 62 63 Returns: 64 A Tensorflow eager tensor 65 """ 66 return pywrap_tfe.TFE_FromDlpackCapsule(dlcapsule, context.context()._handle) 67