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.

[johannahaffner]

Ok, to pick this thing back up -

1. On having just ScalarLike in the documentation, rather than everything it might represent: would it be better to fix the repr for ScalarLike here or by special-casing documentation generation in optimistix?
2. For the Any we get after field conversion of ScalarLike, e.g. in the LearningRate documentation: what exactly do you mean by "better typed wrapper"? And where should that change go? Equinox or optimistix?

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.

3. (I'll keep the equinox decorator that parses defaults that are provided only as arguments to e.g. __init__ in mind, but don't have time to look into that in the coming weeks.)

If it makes more sense to change things in equinox or jaxtyping itself, then I think this can be merged.

[patrick-kidger]

1. I think in jaxtyping!

2. I mean that the current Any is I think coming from introspecting the type hint for jnp.asarray. We could define

   def johannaarray(x: Foo) -> Bar:
       return jnp.asarray(x)
   
   class Foo(eqx.Module):
       x: Bar = eqx.field(converter=johannaarray)

   and I think this would work. (So probably here in Optimistix, next to the module that uses it.)

3. Okay!

[johannahaffner]

1. I'll do it there!
2. Ah! Done, that was easy. Thanks :D

[patrick-kidger]

Awesome! Merged :D