gemini-code-assist[bot] commented on code in PR #347:
URL: https://github.com/apache/tvm-ffi/pull/347#discussion_r2625674765
##########
python/tvm_ffi/cython/device.pxi:
##########
@@ -148,8 +149,12 @@ cdef class Device:
raise ValueError(f"Invalid device index: {parts[1]}")
else:
device_type = device_type_or_name
- if not isinstance(index, int):
- raise TypeError(f"Invalid device index: {index}")
+
+ if not isinstance(index, Integral):
+ if hasattr(index, 'item') and callable(index.item):
+ index = index.item()
+ if not isinstance(index, Integral):
+ raise TypeError(f"Invalid device index: {index}")
Review Comment:
The logic to handle different types for `index` can be simplified for better
readability. The current implementation with nested checks is a bit hard to
follow. A clearer approach is to first attempt to convert `index` using
`.item()` if available, and then perform a single type check. Adding a comment
to explain this conversion would also be beneficial.
```
# Convert torch/numpy scalar to python int
if hasattr(index, 'item') and callable(index.item):
index = index.item()
if not isinstance(index, Integral):
raise TypeError(f"Invalid device index: {index}")
```
##########
python/tvm_ffi/cython/device.pxi:
##########
@@ -129,7 +130,7 @@ cdef class Device:
"trn": DLDeviceType.kDLTrn,
}
- def __init__(self, device_type: str | int, index: Optional[int] = None) ->
None:
+ def __init__(self, device_type: str | int, index: Optional[Integral] =
None) -> None:
Review Comment:
While you've updated the type hint for `index`, the docstring for this
method (lines 79-80) still describes `index` as `int, optional`. Please update
the docstring to reflect that it now also accepts numpy and torch scalars. This
will improve clarity for users of this class.
For example:
```
index : Union[int, numpy.integer, torch.Tensor], optional
Zero-based device index (defaults to ``0`` when omitted). Can be
an integer, a numpy integer scalar, or a single-element torch tensor.
```
--
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]
