Fix filter_and_map docstring argument order for filter_fn and is_leaf_fn

[ayaangazali]

Proposed changes

Module.filter_and_map describes the filter_fn and is_leaf_fn callbacks as taking "a value, the key in which it is found and the containing module", but they are actually called as filter_fn(self, k, v), i.e. (module, key, value). That also matches the type hints (Callable[[Module, str, Any], bool]) and the built-in filters like valid_parameter_filter(module, key, value). The prose just listed the three arguments in reverse order.

Reordered the description to (module, key, value) for both callbacks. Doc only, no behavior change.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

(docstring-only, test_nn.py passes)

---

honest note: freshman contributor here, i use Claude Code to help spot these doc/signature mismatches but i traced the actual call site (filter_fn(self, k, v)) and the parameter filters myself to confirm the order before opening this. let me know if i missed something.