Encodes a megengine tensor to DLPack. Args: tensor (Tensor): The input tensor, and the data type can be `float16`, `float32`, `int8`, `int16`, `int32`, `uint8`, `uint16`, `complex64`. stream (Integer or None): An optional Python integer representing
(tensor, stream=None)
| 7 | |
| 8 | |
| 9 | def to_dlpack(tensor, stream=None): |
| 10 | """ |
| 11 | Encodes a megengine tensor to DLPack. |
| 12 | |
| 13 | Args: |
| 14 | tensor (Tensor): The input tensor, and the data type can be `float16`, `float32`, |
| 15 | `int8`, `int16`, `int32`, `uint8`, `uint16`, `complex64`. |
| 16 | |
| 17 | stream (Integer or None): An optional Python integer representing a CUDA stream, |
| 18 | The current stream is synchronized with this stream before the capsule is created. |
| 19 | If None or -1 is passed then no synchronization is performed. |
| 20 | Returns: |
| 21 | dltensor, and the data type is PyCapsule. |
| 22 | |
| 23 | Examples: |
| 24 | .. code-block:: python |
| 25 | |
| 26 | import megengine as mge |
| 27 | # x is a tensor with shape [2, 3] |
| 28 | x = mge.tensor([[0.2, 0.3, 0.5], |
| 29 | [0.1, 0.2, 0.6]]) |
| 30 | dlpack = mge.utils.dlpack.to_dlpack(x) |
| 31 | print(dlpack) |
| 32 | # <capsule object "dltensor" at 0x7ff04b69cc00> |
| 33 | |
| 34 | """ |
| 35 | |
| 36 | if stream is not None and stream != -1: |
| 37 | return tensor.__dlpack__(stream) |
| 38 | else: |
| 39 | return tensor.__dlpack__() |
| 40 | |
| 41 | |
| 42 | def from_dlpack(ext_tensor, stream=None): |
no test coverage detected