gemini-code-assist[bot] commented on code in PR #97:
URL: https://github.com/apache/tvm-ffi/pull/97#discussion_r2415162794
##########
python/tvm_ffi/stream.py:
##########
@@ -161,3 +161,19 @@ def use_raw_stream(device: core.Device, stream: Union[int,
c_void_p]) -> StreamC
"try use_torch_stream when using torch.cuda.Stream or
torch.cuda.graph"
)
return StreamContext(device, stream)
+
+
+def get_raw_stream(device: core.Device) -> int:
+ """Get the current ffi stream of given device.
+
+ Parameters
+ ----------
+ device : tvm_ffi.Device
+ The device to which the stream belongs.
+
+ Returns
+ -------
+ stream : int
+ The current ffi stream.
+ """
+ return core._env_get_stream(device.dlpack_device_type(), device.index)
Review Comment:
The new function `_env_get_stream` is added to the `core` cython module, but
its signature is missing from the corresponding type stub file
`python/tvm_ffi/core.pyi`. This will cause static type checkers like mypy to
fail when analyzing code that uses `get_raw_stream`.
Please add the function signature to `python/tvm_ffi/core.pyi` to ensure
type safety and maintainability. Based on the existing file, it should be added
after the `_env_set_current_stream` definition.
**Suggested change for `python/tvm_ffi/core.pyi`:**
```python
def _env_get_stream(device_type: int, device_id: int) -> int: ...
```
--
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]
