This is an automated email from the ASF dual-hosted git repository.
cyx-6 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 8e32a61 feat(python): add tensor methods to align with C++ APIs (#604)
8e32a61 is described below
commit 8e32a61eec2e6e0e6d7c8aa1324f22354d5c32ca
Author: Kaining Zhong <[email protected]>
AuthorDate: Thu Jun 4 23:11:09 2026 -0700
feat(python): add tensor methods to align with C++ APIs (#604)
Some of the commonly used tensor methods (`ndim`, `numel`, etc.) are
provided in C++ APIs but not in python APIs. Adding these would make it
handy especially for users who are familiar with pytorch APIs.
---------
Signed-off-by: Kaining Zhong <[email protected]>
---
docs/reference/python/index.rst | 5 ++++-
python/tvm_ffi/core.pyi | 5 +++++
python/tvm_ffi/cython/tensor.pxi | 40 ++++++++++++++++++++++++++++++++++++++++
tests/python/test_tensor.py | 33 +++++++++++++++++++++++++++++++++
4 files changed, 82 insertions(+), 1 deletion(-)
diff --git a/docs/reference/python/index.rst b/docs/reference/python/index.rst
index a59a587..33881c0 100644
--- a/docs/reference/python/index.rst
+++ b/docs/reference/python/index.rst
@@ -48,7 +48,10 @@ Tensor
Device
DLDeviceType
device
-
+ ndim
+ numel
+ size
+ is_contiguous
Function
~~~~~~~~
diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index 892a3b7..c651236 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -160,6 +160,11 @@ class Tensor(Object):
@property
def shape(self) -> tuple[int, ...]: ...
@property
+ def ndim(self) -> int: ...
+ def numel(self) -> int: ...
+ def size(self, idx: int) -> int: ...
+ def is_contiguous(self) -> bool: ...
+ @property
def strides(self) -> tuple[int, ...]: ...
@property
def dtype(self) -> Any: ...
diff --git a/python/tvm_ffi/cython/tensor.pxi b/python/tvm_ffi/cython/tensor.pxi
index dcf8e19..01a9fb0 100644
--- a/python/tvm_ffi/cython/tensor.pxi
+++ b/python/tvm_ffi/cython/tensor.pxi
@@ -276,6 +276,46 @@ cdef class Tensor(CObject):
"""Tensor shape as a tuple of integers."""
return tuple(self.cdltensor.shape[i] for i in
range(self.cdltensor.ndim))
+ @property
+ def ndim(self) -> int:
+ """Number of dimensions of the tensor."""
+ return self.cdltensor.ndim
+
+ def numel(self) -> int:
+ """Total number of elements in the tensor."""
+ cdef int64_t count = 1
+ cdef int i
+ for i in range(self.cdltensor.ndim):
+ count *= self.cdltensor.shape[i]
+ return count
+
+ 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
+ if idx < -ndim or idx >= ndim:
+ raise IndexError(
+ f"Dimension {idx} out of range for tensor with {ndim}
dimensions"
+ )
+ if idx < 0:
+ idx += ndim
+ return self.cdltensor.shape[idx]
+
+ def is_contiguous(self) -> bool:
+ """True if the Tensor is C-contiguous (row-major), False otherwise."""
+ if self.cdltensor.strides == NULL:
+ return True
+ cdef int64_t expected_stride = 1
+ cdef int i
+ cdef int k
+ for i in range(self.cdltensor.ndim, 0, -1):
+ k = i - 1
+ if self.cdltensor.shape[k] == 1:
+ continue
+ if self.cdltensor.strides[k] != expected_stride:
+ return False
+ expected_stride *= self.cdltensor.shape[k]
+ return True
+
@property
def strides(self) -> tuple[int, ...]:
"""Tensor strides as a tuple of integers."""
diff --git a/tests/python/test_tensor.py b/tests/python/test_tensor.py
index cd5ac09..b0bb7c5 100644
--- a/tests/python/test_tensor.py
+++ b/tests/python/test_tensor.py
@@ -39,6 +39,11 @@ 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.ndim == 4
+ assert x.numel() == 640
+ assert x.size(0) == 10
+ assert x.size(-1) == 2
+ assert x.is_contiguous()
assert x.strides == (64, 8, 2, 1)
assert x.dtype == tvm_ffi.dtype("int16")
assert x.device.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLCPU
@@ -47,6 +52,34 @@ def test_tensor_attributes() -> None:
np.testing.assert_equal(x2, data)
+def test_empty_tensor_attributes() -> None:
+ data: npt.NDArray[Any] = np.zeros((4, 0, 4), dtype="int16")
+ if not hasattr(data, "__dlpack__"):
+ return
+ x = tvm_ffi.from_dlpack(data)
+ assert isinstance(x, tvm_ffi.Tensor)
+ assert x.shape == (4, 0, 4)
+ assert x.ndim == 3
+ assert x.strides == (0, 4, 1)
+ assert x.numel() == 0
+ assert x.is_contiguous()
+
+
+def test_non_contiguous_tensor_attributes() -> None:
+ data: npt.NDArray[Any] = np.zeros((4, 4, 4), dtype="int16")
+ slice = data[1:3, :, 1:3]
+ if not hasattr(slice, "__dlpack__"):
+ return
+ x = tvm_ffi.from_dlpack(slice)
+ assert isinstance(x, tvm_ffi.Tensor)
+ assert x.shape == (2, 4, 2)
+ assert x.numel() == 16
+ assert x.size(0) == 2
+ assert x.size(-1) == 2
+ assert not x.is_contiguous()
+ assert x.strides == (16, 4, 1)
+
+
def test_shape_object() -> None:
shape = tvm_ffi.Shape((10, 8, 4, 2))
assert isinstance(shape, tvm_ffi.Shape)
