Adding __init__ arguments to documentation for descents

[johannahaffner]

Hi Patrick,

I would like to add documentation for the __init__ methods of the descents we have.
However, I'm not happy with how this documentation looks like - for instance, we have code that will look something like this

class IndirectDampedNewtonDescent(...):
       root_finder: AbstractRootFinder = Newton(rtol=1e-2, atol=1e-2)
       ...

but when building the documentation, all the default arguments to "Newton" are collected as well and make the documentation much harder to read.

The documentation looks like this, blue is the stuff I'd like removed:

I'd like to write a decorator that is similar to eqxi.docs_remove_args, but which will instead remove all unchanged default values of arguments that are themselves classes, such that the documentation looks more like the code.

Given that the descent is an eqx.Module without an explicit __init__ method, should this be a class decorator?

I'm opening this as a PR here to provide an example. This can be merged once we have the decorator and be a draft for now.

(Name of this branch - this was a place to collect little odds and ends that are not critical and don't need their own PR.)

[johannahaffner]

Opened an issue on equinox: patrick-kidger/equinox#933

[patrick-kidger]

Agreed!

I believe (about 90% sure) that the default arguments are displayed via the default eqx.Module.__repr__.

I think there's a reasonable argument to just always remove default arguments from that display, which lives here:

https://github.com/patrick-kidger/wadler_lindig/blob/1caa67cbc43dfb4e9bebcea7823b7568ce423ee7/wadler_lindig/_definitions.py#L276

[johannahaffner]

Thanks, I will check it out!

Any chance I can fix the type display of converted fields at the same time? ScalarLike gets converted to Any in the documentation of LearningRate, for instance. (This might happen elsewhere, and be less obvious there.)

Code here:

optimistix/optimistix/_solver/learning_rate.py

Line 15 in dcafc48

learning_rate: ScalarLike = eqx.field(converter=jnp.asarray)

[patrick-kidger]

Go for it. The type annotation here is coming from the parameter for jnp.asarray -- probably we just need a better-typed wrapper for it.