Add type annotations for numba and pytorch plugins #5129
Conversation
klecki
force-pushed
the
plugin_function_annotations
branch
from
dd5b04d
to
23e5577
Compare
klecki
changed the title
| class DALIGenericIterator(_DaliBaseIterator): | ||
|
|
||
| def __init__( | ||
| self, | ||
| pipelines: Union[List[Pipeline], Pipeline], | ||
| output_map: List[str], | ||
| size: int = -1, | ||
| reader_name: Optional[str] = None, | ||
| auto_reset: Union[str, bool, None] = False, | ||
| fill_last_batch: Optional[bool] = None, | ||
| dynamic_shape: Optional[bool] = ..., | ||
| last_batch_padded: bool = False, | ||
| last_batch_policy: LastBatchPolicy = LastBatchPolicy.FILL, | ||
| prepare_first_batch: bool = True, | ||
| ) -> None: | ||
| ... | ||
|
|
||
| def __next__(self) -> List[Dict[str, torch.Tensor]]: | ||
| ... |
There was a problem hiding this comment.
I don't think that retyping everything here is the right thing to do - rather the contrary; it's prone to go stale in many subtle ways. Why can't we just annotate the proper code instead?
Even the PEP that defies .pyi files specifies them as a tool rather for injecting type hints than the regular way to do it. I know that we have interface files for automatically generated code or code where the interfaces are mangled in some way (basically all operators). Here it just doesn't apply.
There was a problem hiding this comment.
The technical reason was to keep the TorchPythonFunction without changes, so I put the annotated signatures (that differs from the runtime ones) in the stub file, and thanks to that I need to put all the annotations in the stub, otherwise they become hidden (the annotation file has precedence over which symbols are visible, you can't mix and match this way).
I can move the TorchPythonFunction to a submodule, and move annotations for the iterators back to the file.
There was a problem hiding this comment.
Why do we need a submodule for that?
There was a problem hiding this comment.
So I can have TorchPythonFunction annotated with stub file, as this is needed to expand the definition of **kwargs.
klecki
force-pushed
the
plugin_function_annotations
branch
3 times, most recently
from
2ba7519
to
de70afe
Compare
|
I extracted some adjustments that I found needed for Python function into a separate PR: #5149 |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
!build |
|
CI MESSAGE: [10633692]: BUILD STARTED |
|
CI MESSAGE: [10633692]: BUILD FAILED |
|
!build |
|
CI MESSAGE: [10637571]: BUILD STARTED |
|
CI MESSAGE: [10637571]: BUILD FAILED |
Signed-off-by: Krzysztof Lecki <klecki@nvidia.com> WIP Signed-off-by: Krzysztof Lecki <klecki@nvidia.com> WIP Continue with numba function Signed-off-by: Krzysztof Lecki <klecki@nvidia.com> Numba and python Signed-off-by: Krzysztof Lecki <klecki@nvidia.com> Torch Python Functon Signed-off-by: Krzysztof Lecki <klecki@nvidia.com> Move annotations around Signed-off-by: Krzysztof Lecki <klecki@nvidia.com> Minor revisions Signed-off-by: Krzysztof Lecki <klecki@nvidia.com> Write tests for pytorch, extend TorchFunction overload Fixup the None return in Python Functions (Sequence doesn't count here). Signed-off-by: Krzysztof Lecki <klecki@nvidia.com> Adjust the type of required return type in the PythonFunction callback Signed-off-by: Krzysztof Lecki <klecki@nvidia.com> Extend tests Signed-off-by: Krzysztof Lecki <klecki@nvidia.com>
Signed-off-by: Krzysztof Lecki <klecki@nvidia.com>
Signed-off-by: Krzysztof Lecki <klecki@nvidia.com>
Signed-off-by: Krzysztof Lecki <klecki@nvidia.com>
Signed-off-by: Krzysztof Lecki <klecki@nvidia.com>
Signed-off-by: Krzysztof Lecki <klecki@nvidia.com>
klecki
force-pushed
the
plugin_function_annotations
branch
from
7cef8d7
to
4611987
Compare
|
!build |
|
CI MESSAGE: [10656214]: BUILD STARTED |
|
CI MESSAGE: [10656214]: BUILD PASSED |
Category: New feature
Description:
Type annotations are added for numba and pytorch plugins, covering
the operators (Numba function and Torch Python function) available
in the plugin module.
Additionally, type annotations are added to Pytorch data iterators for DALI.
The code is rearranged so the dictionaries reflect the actual module
hierarchy and allow for inserting the interface .pyi files, rather than
generating the modules at runtime (especially the
fnones).Additional information:
Affected modules and functionalities:
Plugins for Numba and Pytorch (annotations only)
Key points relevant for the review:
Tests:
Checklist
Documentation
DALI team only
Requirements
REQ IDs: N/A
JIRA TASK: N/A