Improved docstring for optax.centralize with explanation and example

[shreyans413]

This PR enhances the docstring for optax.centralize by:

Expanding the explanation of gradient centralization and its benefits.
Clarifying that gradients are centralized along the feature dimension.
Adding an example demonstrating how to apply optax.centralize.

Motivation:
The previous docstring was minimal and lacked details on how gradient centralization works. This update provides a clearer explanation and a practical example, making it easier for users to understand and apply the transformation.

Changes Made:
Updated the docstring with a more detailed description of gradient centralization.
Added a code example demonstrating usage with jax.numpy.
Formatted the reference section properly for better readability.

Checklist:
The docstring follows the project's documentation style.
The example is correctly formatted and functional.
The update does not introduce any breaking changes.

[google-cla]

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.

[vroulet]

I'll take care of these changes internally. Thanks @shreyans413 for the PR

[shreyans413]

Thanks for handling it! I appreciate the review and the opportunity to contribute. Looking forward to contributing again!