NamedTuples not being documented properly

[fabianp]

Note the redundancy and all the "Alias for field number X" in classes such as optax.MultiStepsState: https://optax.readthedocs.io/en/latest/api/optimizer_wrappers.html#optax.MultiStepsState

[amosyou]

I took a look at using sphinx-toolbox to autodoc the namedtuples, but I wasn't able to resolve the duplicate fields showing up. It seems that this is a known issue according to here, not too sure how else to go about fixing this.

I think it's possible to remove the duplicate 'Alias for field ___' but the issue is combining the type with the docstring

[fabianp]

Thanks for looking into it @amosyou . Eliminating the "Alias for field ___" would already be better, even if we loose the type information IMO. Could you post a screenshot of how that would look like?

[amosyou]

this is what it looks like if you remember the :members: line in each of the autoclass directives

[fabianp]

I think that looks good! (module some formatting issues that I thin are orthogonal to this issue). Could you submit a pull request for this?

You're on 🔥 @amosyou !

[amosyou]

revisiting this! so removing :members: from the autoclass directives isn't a viable solution since it also removes fields that are accurately documented like for LookaheadState

i've tried using the autodoc-skip-member event as suggested here but it doesn't seem to be removing the duplicate alias fields. i've also tried modifying the _monkey_patch_doc_strings in conf.py to replace lines with "Alias for field number" with empty string, also didn't work. not sure where to go from here 😅

[vroulet]

Thank you @amosyou for looking into this!
I've just tried removing :members: from the LookAheadState and the fast_state and steps_since_sync were still in the doc. One can even still have types if we are brave enough to make a pass on the whole package.

class LookaheadState(NamedTuple):
  """State of the `GradientTransformation` returned by `lookahead`.

  Attributes:
    fast_state (:class:`optax.OptState`): Optimizer state of the fast optimizer.
    steps_since_sync (``jnp.ndarray``): Number of fast optimizer steps taken since
      slow and fast parameters were synchronized.
  """
  fast_state: base.OptState
  steps_since_sync: jnp.ndarray

[fabianp]

@vroulet taht solutions looks great to me!

[vroulet]

Ok, but I may have missed something, @amosyou can you try again on your side and see if that works? Your idea of removing the :members: sounded great :). I let you see if you want to open a PR to remove the :members: and (if you are extra brave) adding the types.

[amosyou]

oh that looks good! hmm I just tried commenting out :members: again and the entire docstring for LookaheadState is gone from the docs (this has never happened before).

[vroulet]

Strange. All states are missing in that screenshot (they're in pink not in blue). There may be another issue?
Here is the modification of the code on my side:

Flatten
~~~~~~~~
.. autofunction:: flatten

Lookahead
~~~~~~~~~~~~~~~~~
.. autofunction::  lookahead
.. autoclass::  LookaheadParams
   :members:
.. autoclass::  LookaheadState

Masked update
~~~~~~~~~~~~~~
.. autofunction::  masked
.. autoclass::  MaskedState
   :members:

So really just removing the line :members: below LookaheadState

[amosyou]

my bad! commenting in rst is not with # that was why they were missing 🤥

on a side note, I think to have proper rendering of the docstring we need to change Fields to Attributes (even though we're working with namedtuples here since we're using autoclass). hope that's fine?

i'll create a PR for this. thanks @vroulet!

[vroulet]

Yes, typically, MultiStepsState should be reformatted with attributes instead of fields. Thanks again for looking into that!

[amosyou]

also for ApplyIfFiniteState, I noticed that all the fields are typed as Any. Did we figure out a solution for the typing jnp.array?

[vroulet]

also for ApplyIfFiniteState, I noticed that all the fields are typed as Any. Did we figure out a solution for the typing jnp.array?

They should be Union[jax.Array, int] for notfinite_count, last_finite, and total_notfinite, I believe.
For inner_state it should be the usual base.Optstate

[vroulet]

my bad! commenting in rst is not with # that was why they were missing 🤥

on a side note, I think to have proper rendering of the docstring we need to change Fields to Attributes (even though we're working with namedtuples here since we're using autoclass). hope that's fine?

i'll create a PR for this. thanks @vroulet!

PS: no need to add types for all attributes (it's possible but it would take forever). Simply remove the :members: in the doc. If people want to know the types, they should be able to click on [source] in the doc.
That said, if you find some issues (like the fields in MultiStepsState (which is badly formatted anyway) or like in the ApplyFiniteState you pointed out), would be super cool if you could fix them by the way. We can review that in your PR together.