add comment
This commit is contained in:
parent
29856e6354
commit
dab3c26d03
@ -22,12 +22,41 @@ from tensorflow.python import pywrap_tfe
|
|||||||
from tensorflow.python.util.tf_export import tf_export
|
from tensorflow.python.util.tf_export import tf_export
|
||||||
|
|
||||||
|
|
||||||
# tf.dlpack.to_dlpack/from_dlpack doesn't work. How to fix?
|
|
||||||
@tf_export("experimental.dlpack.to_dlpack", v1=[])
|
@tf_export("experimental.dlpack.to_dlpack", v1=[])
|
||||||
def to_dlpack(tf_tensor):
|
def to_dlpack(tf_tensor):
|
||||||
|
"""Returns the dlpack capsule representing the tensor. This operation
|
||||||
|
ensures the underlying data memory is ready when returns.
|
||||||
|
|
||||||
|
```python
|
||||||
|
a = tf.tensor([1, 10])
|
||||||
|
dlcapsule = tf.experimental.dlpack.to_dlpack(a)
|
||||||
|
# dlcapsule represents the dlpack data structure
|
||||||
|
```
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tf_tensor: Tensorflow eager tensor, to be converted to dlpack capsule.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A PyCapsule named as dltensor, which shares the underlying memory to other
|
||||||
|
framework. This PyCapsule can be consumed only once.
|
||||||
|
"""
|
||||||
return pywrap_tfe.TFE_ToDlpackCapsule(tf_tensor)
|
return pywrap_tfe.TFE_ToDlpackCapsule(tf_tensor)
|
||||||
|
|
||||||
|
|
||||||
@tf_export("experimental.dlpack.from_dlpack", v1=[])
|
@tf_export("experimental.dlpack.from_dlpack", v1=[])
|
||||||
def from_dlpack(dlcapsule):
|
def from_dlpack(dlcapsule):
|
||||||
|
"""Returns the Tensorflow eager tensor, which uses the memory shared by
|
||||||
|
dlpack capsules from other framework.
|
||||||
|
|
||||||
|
```python
|
||||||
|
a = tf.experimental.dlpack.from_dlpack(dlcapsule)
|
||||||
|
# `a` uses the memory shared by dlpack
|
||||||
|
```
|
||||||
|
|
||||||
|
Args:
|
||||||
|
dlcapsule: A PyCapsule named as dltensor
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A Tensorflow eager tensor
|
||||||
|
"""
|
||||||
return pywrap_tfe.TFE_FromDlpackCapsule(dlcapsule)
|
return pywrap_tfe.TFE_FromDlpackCapsule(dlcapsule)
|
||||||
|
Loading…
Reference in New Issue
Block a user