Add ACProp optimizer.

[carlosgmartin]

#965

[carlosgmartin]

@fabianp @vroulet This PR is getting the same error as #962. Looks like there's a problem with optax/monte_carlo/control_variates_test.py > ConsistencyWithStandardEstimators.testWeightedLinearFunction > _score_function_jacobians & control_delta_method.

[vroulet]

Yep, we know. This bug does not show up internally, neither on a mac, just on linux (which I don't have).
It's related to the new jax release. We are working on it.

[fabianp]

@carlosgmartin : we could also use some help if you have cycles to look into this ...

[vroulet]

It's been fixed in #972. Should be merged tomorrow.

[vroulet]

Thanks one more time for this contribution @carlosgmartin
This looks very good.
One additional comment: since this paper is cited only 6 times, we will place it in the contrib folder for now, see https://optax.readthedocs.io/en/latest/development.html#inclusion-criteria
(that's a policy implemented only recently)

[carlosgmartin]

@vroulet Done.

[vroulet]

Thanks! Can you add it to the list of common tests of the contrib folder?

[carlosgmartin]

@vroulet Done.

[vroulet]

Perfect, thank you again @carlosgmartin.

• The scale_by_sophia can be added to the doc.
• You may format the reference in scale_by_sophia.
• The copyright should be 2024.

I can take it from here if you don't have time.
We sincerely appreciate the dedication in your PRs.

[carlosgmartin]

@vroulet Did you mean adding scale_by_acprop to docs/api/contrib.rst in your first bullet point?

[vroulet]

Yes. People may want to combine this scale function with e.g. a clipping, so the documentation helps.
This is just a nitpick to finalize the contrib.

[vroulet]

Perfect, I'll take it from here (I need to make some internal changes to add this optimizer to our builds).
Thanks again