search

v0lta / PyTorch-Wavelet-Toolbox

Differentiable fast wavelet transforms in PyTorch with GPU support.
https://pytorch-wavelet-toolbox.readthedocs.io
European Union Public License 1.2
279 stars 36 forks source link

Type Hints #85

Closed fzimmermann89 closed 3 months ago

fzimmermann89 commented 4 months ago

Hi,

thank you for providing this toolbox. We noticed some problems with the type hints:

waverec(coeffs: List[torch.Tensor]) -> It would be nice to use Sequence[torch.Tensor] here. This would also allow to be called with a tuple of Tensors. The function should not modify the list of tensors nor depend on it actually being a list.

wavedec2( ... ) -> List[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]] It would be nicer if this could be changed to return tuple[torch.Tensor, tuple[*tuple[torch.Tensor, torch.Tensor, torch.Tensor],...] This would mean that the return is always a tuple with the first entry a tensor, followed by an arbitrary number of tuples of 3 tensors. Currently, the type hint says that a) the first entry can also be a tuple and b) the later entries can also be tensors (although they should only be tuples..)

It would be nice to return tighter types. And take sequences instead of lists as input for all public functions.

I would be willing to create a PR if there would be interest do do the change.

v0lta commented 4 months ago

Yes, this sounds good. I would be happy to review a PR.

v0lta commented 3 months ago

@felixblanke 's pull request fixes the issue.