Skip to content

docs(nnx): fix training API reference inconsistencies - #5544

Open
greekera1000 wants to merge 2 commits into
google:mainfrom
greekera1000:docs/nnx-training-api-reference
Open

docs(nnx): fix training API reference inconsistencies#5544
greekera1000 wants to merge 2 commits into
google:mainfrom
greekera1000:docs/nnx-training-api-reference

Conversation

@greekera1000

@greekera1000 greekera1000 commented Aug 4, 2026

Copy link
Copy Markdown

What does this PR do?

Fixes a few inconsistencies in the flax.nnx.training API reference (part of #5161):

  • EMA.apply_to() is public and used in EMA's own docstring examples, but wasn't rendered in the reference — added it to :members:.
  • Accuracy listed only update while the sibling metrics list __init__, reset, update, compute — aligned it (this also surfaces the threshold constructor doc).
  • Statistics (returned by Welford.compute()) had no docstring — added one.

Scope: docs + docstrings only, no behavioral changes. All training module doctests pass locally.

Relates to #5161.

Checklist

  • This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).
  • This change is discussed in a Github issue: [docs]: Improve docs for NNX API Reference #5161
  • The documentation and docstrings adhere to the documentation guidelines.
  • This change includes necessary high-coverage tests. (No quality testing = no merge!)

@google-cla

google-cla Bot commented Aug 4, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@greekera1000

Copy link
Copy Markdown
Author

@googlebot I signed it!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant