Feature Request: Add default parameter values to docstring

[jenkspt]

Currently the docstrings for Module subclasses don't display the default argument values when using ?. Not sure if there is an automated way to display this, or to update all of the docstrings.

[patrick-kidger]

I'm not sure what you mean about ? - I assume this is in reference to some piece of tooling you use?

In any case, I'd be happy to include default parameters in docstrings. It may be possible to do this in automated fashion either now or in the future. See this Reddit thread.

[jenkspt]

In IPython and Jupyter, a function/class or object followed by ? will display the signature and docstring. i.e. np.split?

Signature: np.split(ary, indices_or_sections, axis=0)
Docstring:
Split an array into multiple sub-arrays as views into `ary`.

Parameters
----------
...

I use this so much that I didn't realize it isn't standard python. help in the python shell is very similar.

[paganpasta]

Adding my experience to this issue. Using an IDE (pycharm) I am able to get suggestions for the input parameters for most of the subclasses. But, nn.Sequential fails to parse due to lack of __init__ and gets highlighted as a warning blob.

[jenkspt]

After some more investigation: For some reason eqx.Module subclasses don't use the signature from the __init__ method. For example:

class Test:
    def __init__(a: int, b:bool=False):
        pass

help(Test) shows:

class Test(builtins.object)
 |  Test(a: bool, b: int = 1)
 |  
 |  Methods defined here:
 |  
 |  __init__(self, a: bool, b: int = 1)
 |      Initialize self.  See help(type(self)) for accurate signature.
...

Here the init signature Test(a: bool, b: int = 1) is grabbed from the __init__ method.

However with eqx.Module subclasses:

`help(nn.Linear)```

class Linear(equinox.module.Module)
 |  Linear(*args, **kwargs)
 |  
 |  Performs a linear transformation.
 |  
 |  Method resolution order:
 |      Linear
 |      equinox.module.Module
 |      builtins.object
 |  
 |  Methods defined here:
...
 |  __init__(self, in_features: int, out_features: int, use_bias: bool = True, *, key: 'jax.random.PRNGKey')
 |      **Arguments:**
 |      
 |      - `in_features`: The input size.
 |      - `out_features`: The output size.
 |      - `use_bias`: Whether to add on a bias as well.
 |      - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
 |          initialisation. (Keyword only argument.)
...

Here the init signature is Linear(*args, **kwargs). I'm assuming this has something to do with the way _ModuleMeta works.

I found a hack to fix this is to just add cls.__signature__ = inspect.signature(cls.__init__) to the _ModuleMeta.__new__ method

Not sure if there are any downsides to doing this -- and maybe there is another better way to fix this in the metaclass.

[patrick-kidger]

Ah, this will probably be because it's grabbing the metaclass __call__ rather than using the class __init__. (I suppose help must have some special-casing internally to not exhibit the same behaviour for type.__call__.)

I think overriding cls.__signature__ is probably the correct thing to do. I'd be happy to accept a PR fixing this.

[patrick-kidger]

Closing with #573, as I think this is now resolved.