Skip to content

docs(nnx): add missing functional args to Conv and LinearGeneral #5174

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
ayulockedin wants to merge 7 commits into
base: main
Choose a base branch
from
Open

docs(nnx): add missing functional args to Conv and LinearGeneral #5174

ayulockedin wants to merge 7 commits into from
+5 −0

Conversation

@ayulockedin
Copy link
Contributor

@ayulockedin ayulockedin commented Jan 6, 2026 •
edited
Loading

What does this PR do?

Description

This PR updates the docstrings for Conv and LinearGeneral to include functional arguments that were present in __init__ but missing from the API reference.

Changes

  • LinearGeneral: Added dot_general and dot_general_cls to the Args list.
  • Conv: Added conv_general_dilated to the Args list.

Related Issues

  • Contributes to the documentation cleanup initiative.

Great, you are contributing to Flax!

Fixes # (issue)

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/discussion (please add a link).
  • The documentation and docstrings adhere to the documentation guidelines.
  • This change includes necessary high-coverage tests. (No quality testing = no merge!)

@ayulockedin
Copy link
Contributor Author

ayulockedin commented Jan 6, 2026

Hi @cgarciae , following up on the doc cleanups: this PR fills in some missing arguments in the Conv and LinearGeneral docstrings. Ready for review when you have a moment!

vfdev-5
vfdev-5 reviewed Jan 7, 2026
View reviewed changes
Copy link
Collaborator

@vfdev-5 vfdev-5 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR @ayulockedin
I left few comments to improve the docs a bit more.

flax/nnx/nn/linear.py Outdated
param_dtype: the dtype passed to parameter initializers (default: float32).
kernel_init: initializer function for the weight matrix.
bias_init: initializer function for the bias.
dot_general: dot product function (default: None).
Copy link
Collaborator

@vfdev-5 vfdev-5 Jan 7, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a note about if none of dot_general or dot_general_cls are provided jax.lax.dot_general is used.

Copy link
Contributor Author

@ayulockedin ayulockedin Jan 7, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done!!

ayulockedin
ayulockedin commented Jan 7, 2026
View reviewed changes
Copy link
Contributor Author

@ayulockedin ayulockedin left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed the Conv layer docstring that is failing the tests (likely due to line length)

@ayulockedin
Copy link
Contributor Author

ayulockedin commented Jan 7, 2026

@vfdev-5 Thanks for the guidance! I've updated the docstrings to clarify the dot_general fallback and fixed the line wrapping in the Conv layer. All checks have passed. Please let me know if everything looks good!

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

Reviewers

@vfdev-5 vfdev-5 vfdev-5 left review comments

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@ayulockedin @vfdev-5