gemini-code-assist[bot] commented on code in PR #604:
URL: https://github.com/apache/tvm-ffi/pull/604#discussion_r3352491832
##########
python/tvm_ffi/cython/tensor.pxi:
##########
@@ -276,6 +276,52 @@ 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
Review Comment:
For empty tensors (where any dimension is 0, e.g., shape `(3, 0, 4)`), the
current implementation of `is_contiguous` will incorrectly return `False`. This
is because after encountering the 0-sized dimension, `expected_stride` becomes
`0`, causing subsequent stride checks for non-zero dimensions to fail (e.g.,
`strides[0] != 0`). Empty tensors are always considered contiguous by standard
convention (e.g., in NumPy and PyTorch). Adding an early return when
`self.numel() == 0` resolves this issue.
```
def is_contiguous(self) -> bool:
"""True if the Tensor is C-contiguous (row-major), False
otherwise."""
if self.numel() == 0:
return True
if self.cdltensor.strides == NULL:
return True
```
##########
tests/python/test_tensor.py:
##########
@@ -39,6 +39,12 @@ 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()
Review Comment:
It is highly recommended to also test `is_contiguous()` on non-contiguous
tensors (e.g., transposed arrays) and empty tensors with 0-sized dimensions
(e.g., shape `(3, 0, 4)`) to ensure correctness across all layouts.
```suggestion
assert x.is_contiguous()
non_contig_data = data.T
if hasattr(non_contig_data, "__dlpack__"):
y = tvm_ffi.from_dlpack(non_contig_data)
assert not y.is_contiguous()
empty_data = np.zeros((3, 0, 4), dtype="int16")
if hasattr(empty_data, "__dlpack__"):
z = tvm_ffi.from_dlpack(empty_data)
assert z.is_contiguous()
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
