fix(contrib): correct momo and momo_adam docstring examples to use apply_updates

[wahid18-maqs]

Summary

This PR corrects the usage examples in the momo() and momo_adam() docstrings to follow the standard Optax optimizer update pattern and adds doctest coverage to keep the examples executable.

Changes

• Update the momo() and momo_adam() docstring examples to:
  • capture updates from solver.update(...),
  • apply them using optax.apply_updates(...), and
  • update the expected outputs to match the corrected behavior.
• Add a dedicated doctest wrapper in optax/contrib/_momo_test.py to execute the examples as part of the test suite.

Why

The previous examples assigned the return value of solver.update() directly to params, even though update() returns (updates, state). Users copying the examples would end up treating updates as parameters, leading to incorrect optimization behavior.

Adding doctest coverage also ensures the examples remain valid and synchronized with the implementation.

The doctest suite passes successfully