Added clarification to apply_fn parameter of apply_to_modules

[yuchaoran2011]

Clarified the apply_fn argument in the docstring to include details about its parameters.

Proposed changes

When creating a transformer architecture using MLX and working on weight initialization, I had to read the source code of apply_to_modules to understand what apply_fn expects as its input parameters. Modified the doctoring to add clarification in the documentation.

Checklist

Put an x in the boxes that apply.

• [ x] I have read the CONTRIBUTING document
• [ x] 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
• [ x] I have updated the necessary documentation (if needed)

[awni]

Ty!