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 90dba57 fix(Array): Make array slicing behavior to return list (#50)
90dba57 is described below
commit 90dba57cf810e7fb6a5ad8316bcf4c53c699d51c
Author: Junru Shao <[email protected]>
AuthorDate: Tue Sep 23 16:01:42 2025 -0700
fix(Array): Make array slicing behavior to return list (#50)
This PR fixes a behavioral inconsistency between existing TVM stack and
the latest `main`.
Originally, slicing a `tvm_ffi.Array[T]` always returns `list[T]`. This
behavior was altered by PR #37, which instead returns `Array[T]`,
breaking a few unittests.
---
python/tvm_ffi/container.py | 8 +++-----
tests/python/test_container.py | 1 +
2 files changed, 4 insertions(+), 5 deletions(-)
diff --git a/python/tvm_ffi/container.py b/python/tvm_ffi/container.py
index f179fc9..d4aa24f 100644
--- a/python/tvm_ffi/container.py
+++ b/python/tvm_ffi/container.py
@@ -117,14 +117,12 @@ class Array(core.Object, Sequence[T]):
def __getitem__(self, idx: SupportsIndex, /) -> T: ...
@overload
- def __getitem__(self, idx: slice, /) -> Array[T]: ...
+ def __getitem__(self, idx: slice, /) -> list[T]: ...
- def __getitem__(self, idx: SupportsIndex | slice, /) -> T | Array[T]:
- """Return one element or a new :class:`Array` for a slice."""
+ def __getitem__(self, idx: SupportsIndex | slice, /) -> T | list[T]:
+ """Return one element or a list for a slice."""
length = len(self)
result = getitem_helper(self, _ffi_api.ArrayGetItem, length, idx)
- if isinstance(result, list):
- return cast(Array[T], type(self)(result))
return result
def __len__(self) -> int:
diff --git a/tests/python/test_container.py b/tests/python/test_container.py
index 54b41b7..3f3465e 100644
--- a/tests/python/test_container.py
+++ b/tests/python/test_container.py
@@ -28,6 +28,7 @@ def test_array() -> None:
assert len(a) == 3
assert a[-1] == 3
a_slice = a[-3:-1]
+ assert isinstance(a_slice, list) # TVM array slicing returns a list[T]
instead of Array[T]
assert (a_slice[0], a_slice[1]) == (1, 2)
