This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 4108fd22 [FEAT] Expose Tensor data_ptr in Python (#679)
4108fd22 is described below
commit 4108fd226e0048d0e5bfd53b7fc3434e67b03052
Author: Yixin Dong <[email protected]>
AuthorDate: Wed Jul 22 15:19:03 2026 -0700
[FEAT] Expose Tensor data_ptr in Python (#679)
## Summary
- add `Tensor.data_ptr()` to return the address of the first tensor
element
- include the DLPack tensor's byte offset in the returned address
- add the Python type stub and tests using real NumPy arrays and sliced
views
This allows callers to access tensor memory directly without exporting
and parsing a DLPack capsule.
## Test plan
- [x] `uv run pytest tests/python/test_tensor.py -q` (8 passed, 3
skipped)
- [x] `uv run pytest tests/python -q` (2418 passed, 41 skipped, 2
expected failures)
- [x] `uv run pre-commit run --files python/tvm_ffi/cython/base.pxi
python/tvm_ffi/cython/tensor.pxi python/tvm_ffi/core.pyi
tests/python/test_tensor.py`
---
python/tvm_ffi/core.pyi | 1 +
python/tvm_ffi/cython/base.pxi | 2 +-
python/tvm_ffi/cython/tensor.pxi | 10 ++++++++++
tests/python/test_tensor.py | 10 ++++++++++
4 files changed, 22 insertions(+), 1 deletion(-)
diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index 525d7a40..29659bd6 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -169,6 +169,7 @@ class Tensor(Object):
@property
def ndim(self) -> int: ...
def numel(self) -> int: ...
+ def data_ptr(self) -> int: ...
def size(self, idx: int) -> int: ...
def is_contiguous(self) -> bool: ...
@property
diff --git a/python/tvm_ffi/cython/base.pxi b/python/tvm_ffi/cython/base.pxi
index 56c4d3d1..20b9306a 100644
--- a/python/tvm_ffi/cython/base.pxi
+++ b/python/tvm_ffi/cython/base.pxi
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
import ctypes
-from libc.stdint cimport int32_t, int64_t, uint64_t, uint32_t, uint8_t, int16_t
+from libc.stdint cimport int32_t, int64_t, uint64_t, uint32_t, uint8_t,
int16_t, uintptr_t
from libc.string cimport memcpy
from libcpp.vector cimport vector
from cpython.bytes cimport PyBytes_AsStringAndSize, PyBytes_FromStringAndSize,
PyBytes_AsString
diff --git a/python/tvm_ffi/cython/tensor.pxi b/python/tvm_ffi/cython/tensor.pxi
index 21b03108..e267d9c8 100644
--- a/python/tvm_ffi/cython/tensor.pxi
+++ b/python/tvm_ffi/cython/tensor.pxi
@@ -289,6 +289,16 @@ cdef class Tensor(CObject):
count *= self.cdltensor.shape[i]
return count
+ def data_ptr(self) -> int:
+ """Return the address of the first tensor element.
+
+ Returns
+ -------
+ data_ptr
+ Address of the first tensor element, including the tensor's byte
offset.
+ """
+ return <uintptr_t>self.cdltensor.data + self.cdltensor.byte_offset
+
def size(self, idx: int) -> int:
"""Get the size of the ``idx``-th dimension. Negative ``idx`` counts
from the last dimension."""
cdef int ndim = self.cdltensor.ndim
diff --git a/tests/python/test_tensor.py b/tests/python/test_tensor.py
index e21de826..f9a559aa 100644
--- a/tests/python/test_tensor.py
+++ b/tests/python/test_tensor.py
@@ -53,6 +53,16 @@ def test_tensor_attributes() -> None:
np.testing.assert_equal(x2, data)
+def test_tensor_data_ptr() -> None:
+ data = np.arange(8, dtype="int32")
+ tensor = tvm_ffi.from_dlpack(data)
+ assert tensor.data_ptr() == data.ctypes.data
+
+ view = data[2:]
+ tensor_view = tvm_ffi.from_dlpack(view)
+ assert tensor_view.data_ptr() == view.ctypes.data
+
+
def test_empty_tensor_is_contiguous() -> None:
# Empty tensors are trivially contiguous regardless of what
# strides the producer reports (numpy 2.3+ via __dlpack__ now
