ENH: uarray support for scipy.signal

[Smit-create]

Reference issue: scipy#14353

Reference PR from which I followed the implementation: scipy#14356

cc @rgommers

[rgommers]

@Smit-create there are a lot of commits in this PR that don't belong here. My master branch is up-to-date with upstream master right now. Can you rebase to get rid of the commits that are unrelated to uarray for scipy.signal?

@AnirudhDagar maybe you'd be able to do a first review on this PR? It's still small - Uarray support for a single function right now.

@Smit-create I'll be travelling for the next couple of days, so I don't think I'll be able to look at this before Thursday.

[Smit-create]

Can you rebase to get rid of the commits that are unrelated to uarray for scipy.signal?

Done!

I'll be travelling for the next couple of days, so I don't think I'll be able to look at this before Thursday.

Okay, till then I will have a look at some sample code in colab for verifying this PR.

[Smit-create]

@AnirudhDagar Once you review it in your free time let me know if I can proceed over this or I need to add something more.

[Smit-create]

I have created a sample cusignal_backend for testing in colab.

[AnirudhDagar]

Thanks @Smit-create for your PR. This is already looking like a great start!

I have a few thoughts that might help once you add more functions. Given the number of scipy.signal submodules and the different signatures they have, you will probably need to adapt the argument replacers accordingly and define a few more than just the _h_x_replacer (similar to ndimage). That said, this can get a bit out of hand and testing becomes slightly tedious if we don't use __ua_convert__. It's great that you have created a cusignal backend for testing (that's how I started testing ndimage in cupy as well) but we should also focus on SciPy CI (we can't add cusignal backends there), so, until we actually use __ua_convert__ within the SciPy backend, the mock backend tests on their own don't really mean much in terms of testing the scipy.signal api. With __ua_convert__ we can leverage the scipy.signal test suite already in place. This is discussed in scipy#14356, as the overhead is not too much for __ua_convert__. Maybe we should incorporate it here as well once you start adding more functions.

On a different note, about documentation of functions (I'm assuming you have left that out intentionally given the PR is still WIP).
scipy/signal/__init__.py will need some documentation on those uarray backend control functions (See this example). Although on a second thought I'm not sure if it is a good idea to duplicate the documentation for these functions and in fact defining these functions again for each and every module that supports uarray since the functions still remain the same. It is still important to add some documentation explaining to the users that the signal module is uarray supported and then maybe we document all these 4 common functions under the uarray.rst doc page? What do you think @Smit-create? I remember @czgdp1807 brought this up some time ago.

I'm yet to test the cusignal backend that you shared but I made a cursory look, seems good to me and you can probably go ahead adding more methods. :)

[czgdp1807]

Not sure fully, but these functions might get repeated across modules. We can factor out the common parts later though (dealing with concrete things is better than doing abstract thinking way before things get less hazy).

[AnirudhDagar]

Yes, I feel the same way. I remember you mentioned that in scipy#14407 and I guess this comment explains my thoughts on the approach since we still need different documentation.

[AnirudhDagar]

Maybe once all the uarray PRs land (possibly starting with scipy#14356), we can work on refactoring the common parts into a signal place which can then be used in all the modules.

[czgdp1807]

in those functions which can also be addressed through some utility doc functions or we can completely drop module-specific examples and make these functions generalized.

I see. Well may be we can keep the module specific parts (like Examples section) in a common place. I don't see any harm in doing that. Like give examples from 2-3 modules in doc-string of common function (say set_backend).

[czgdp1807]

we can work on refactoring the common parts into a signal place which can then be used in all the modules.

Yes, that would be way more easier than thinking of possible cases now. We don't know how things might end up in distant future.

[Smit-create]

Maybe once all the uarray PRs land (possibly starting with scipy#14356), we can work on refactoring the common parts into a signal place which can then be used in all the modules.

+1, refactoring seems an easy option later as we will be more clear which all share common behaviour.

[czgdp1807]

I remember @czgdp1807 brought this up some time ago.

Are you talking about this function for duplicating docs? Is there any major drawback of duplicating the docs from regular SciPy functions to multimethods? Well, I think it's redundant and introduces a slight overhead (copy pasting doc strings from one variable to another, though no objective data to support this). Other than that I can't think of anything else.

[AnirudhDagar]

Are you talking about this function for duplicating docs?

No not this, although I'm aware this can sit better in some common scipy utils place as well like doccer.py. I was actually talking about the backend control functions (set_global_backend , register_backend,.. ) and their documentation like you mentioned in your comment above.

[Smit-create]

Thanks for the review @AnirudhDagar and @czgdp1807

probably need to adapt the argument replacers accordingly and define a few more than just the _h_x_replacer

Yes, will do that in further commits.

until we actually use __ua_convert__ within the SciPy backend, the mock backend tests on their own don't really mean much in terms of testing the scipy.signal api. With __ua_convert__ we can leverage the scipy.signal test suite already in place.

Makes sense to me.

It is still important to add some documentation explaining to the users that the signal module is uarray supported and then maybe we document all these 4 common functions under the uarray.rst doc page?

I think once we complete one complete module, we could add some docs that can show how to use the backends (but again should we wait until we add backends/*_backend.py in the modules that support uarray?)

[Smit-create]

@rgommers @AnirudhDagar
I created cupy_backend on the other branch and installed it on the colab notebook. It seems to be working fine. Please have a look at the notebook and feel free to provide comments over there too.

[Smit-create]

There seems to be some issue with all_of_type:

>>> from scipy._lib.uarray import all_of_type
>>> import numpy as np
>>> @all_of_type(np.ndarray)
... def f(p):
...     return p
... 
>>> sos_f32  = np.array([[4., 5., 6., 1., 2., 3.]], dtype=np.float32)
>>> sos_f32
array([[4., 5., 6., 1., 2., 3.]], dtype=float32)
>>> sos_f32.shape
(1, 6)
>>> f(sos_f32)
(<Dispatchable: type=<class 'numpy.ndarray'>, value=array([4., 5., 6., 1., 2., 3.], dtype=float32)>,)
>>> f(sos_f32)[0].value.shape
(6,)

While I was trying to add __ua_convert__, there were many test failures due to changes in the dimensions of numpy arrays. I was expecting sos_f32 and f(sos_f32)[0].value of same dimensions

[rgommers]

I think all_of_type doesn't work for multimethod definitions that just return their input - which seems wrong, a 1-element tuple needs to be returned. The for arg in extracted_args in this all_of_type body:

            extracted_args = tuple(func(*args, **kwargs))
            return tuple(
                Dispatchable(arg, arg_type)
                if not isinstance(arg, Dispatchable)
                else arg
                for arg in extracted_args
            )

is assuming extracted_args is a tuple; if it's a single 2-D array then it's going to unpack it into 1-D arrays.

I'm not sure if it's a bug that needs fixing (an assert may be useful), or if the docs need to be updated. Either way, all_of_type isn't tested in uarray, so let's add a few tests and improve the docs. Can you open an issue on the Uarray repo with a reproducer @Smit-create?

This fixes things

@all_of_type(np.ndarray)
def f(p):
    return (p,)

# OR:
def f(p):
    return (Dispatchable(p, np.ndarray),)

[Smit-create]

With the present commits, __ua_convert__ works fine for the added functions. So, next, I'll look into each function's parameters (optional arguments) which might also need to be dispatched.

[czgdp1807]

Can we use a code generator to produce _multimethods.py file? Like a yml file which can contain all the information. We just need to define the replacer methods and these kind of 3-4 liner blocks can be generated dynamically? cc: @grlee77 @rgommers @AnirudhDagar

P.S. It might not be worth it because anyways its the same amount of code in both cases (auto generation and manually writing). But still thought to put this up.

[AnirudhDagar]

I know that sounds appealing and probably better engineered (will help with consistency) but I believe there are not many uarray modules and multimethods left to be added now (with the PRs already open) and autogeneration might just make things more complicated (one more thing to be handled).

[Smit-create]

@rgommers @czgdp1807 @AnirudhDagar This PR now looks good for a first thorough review.