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 8377011 [CYTHON] Expose Tensor.strides (#102)
8377011 is described below
commit 83770118612208d011032a9e17ca93a3bd3d85e4
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 12 12:56:03 2025 -0400
[CYTHON] Expose Tensor.strides (#102)
This PR updates the Tensor class to expose strides, also bump DLPack to
v1.2
---
3rdparty/dlpack | 2 +-
pyproject.toml | 2 +-
python/tvm_ffi/__init__.py | 2 +-
python/tvm_ffi/core.pyi | 3 +++
python/tvm_ffi/cython/tensor.pxi | 19 +++++++++++++++++++
tests/python/test_tensor.py | 1 +
6 files changed, 26 insertions(+), 3 deletions(-)
diff --git a/3rdparty/dlpack b/3rdparty/dlpack
index 1117366..93c8f2a 160000
--- a/3rdparty/dlpack
+++ b/3rdparty/dlpack
@@ -1 +1 @@
-Subproject commit 111736618e8d1028b23605f76dcaa6a38cfea809
+Subproject commit 93c8f2a3c774b84af6f652b1992c48164fae60fc
diff --git a/pyproject.toml b/pyproject.toml
index 25a8137..a0a00c9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,7 +17,7 @@
[project]
name = "apache-tvm-ffi"
-version = "0.1.0b15"
+version = "0.1.0b16"
description = "tvm ffi"
authors = [{ name = "TVM FFI team" }]
diff --git a/python/tvm_ffi/__init__.py b/python/tvm_ffi/__init__.py
index e46a7aa..6ddf32f 100644
--- a/python/tvm_ffi/__init__.py
+++ b/python/tvm_ffi/__init__.py
@@ -17,7 +17,7 @@
"""TVM FFI Python package."""
# version
-__version__ = "0.1.0b15"
+__version__ = "0.1.0b16"
# order matters here so we need to skip isort here
# isort: skip_file
diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index 05e5a04..3496ece 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -410,6 +410,9 @@ class Tensor(Object):
def shape(self) -> tuple[int, ...]:
"""Tensor shape as a tuple of integers."""
@property
+ def strides(self) -> tuple[int, ...]:
+ """Tensor strides as a tuple of integers."""
+ @property
def dtype(self) -> Any:
"""Data type as :class:`tvm_ffi.dtype` (``str`` subclass)."""
@property
diff --git a/python/tvm_ffi/cython/tensor.pxi b/python/tvm_ffi/cython/tensor.pxi
index 4ebc515..acdb4fd 100644
--- a/python/tvm_ffi/cython/tensor.pxi
+++ b/python/tvm_ffi/cython/tensor.pxi
@@ -179,6 +179,18 @@ def _shape_obj_get_py_tuple(obj):
return tuple(shape.data[i] for i in range(shape.size))
+def _make_strides_from_shape(tuple shape):
+ cdef int64_t expected_stride = 1
+ cdef list strides = []
+ cdef int64_t ndim = len(shape)
+ cdef int64_t reverse_index
+ for i in range(ndim):
+ reverse_index = ndim - i - 1
+ strides.append(expected_stride)
+ expected_stride *= shape[reverse_index]
+ return tuple(reversed(strides))
+
+
cdef class Tensor(Object):
"""Tensor object that represents a managed n-dimensional array.
"""
@@ -189,6 +201,13 @@ cdef class Tensor(Object):
"""Shape of this array"""
return tuple(self.cdltensor.shape[i] for i in
range(self.cdltensor.ndim))
+ @property
+ def strides(self):
+ """Strides of this array"""
+ if self.cdltensor.strides == NULL:
+ return _make_strides_from_shape(self.shape)
+ return tuple(self.cdltensor.strides[i] for i in
range(self.cdltensor.ndim))
+
@property
def dtype(self):
"""Data type of this array"""
diff --git a/tests/python/test_tensor.py b/tests/python/test_tensor.py
index 186d91b..9b73ce9 100644
--- a/tests/python/test_tensor.py
+++ b/tests/python/test_tensor.py
@@ -38,6 +38,7 @@ def test_tensor_attributes() -> None:
x = tvm_ffi.from_dlpack(data)
assert isinstance(x, tvm_ffi.Tensor)
assert x.shape == (10, 8, 4, 2)
+ assert x.strides == (64, 8, 2, 1)
assert x.dtype == tvm_ffi.dtype("int16")
assert x.device.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLCPU
assert x.device.index == 0
