Add ACProp optimizer.

[carlosgmartin]

#965

[carlosgmartin]

@fabianp @vroulet This PR is getting the same error as https://github.com/google-deepmind/optax/pull/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.

[carlosgmartin]

@vroulet Done.

[vroulet]

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

[carlosgmartin]

@vroulet Done.

[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

[carlosgmartin]

@vroulet Should I close this now?

[vroulet]

It should have been closed automatically but yes, let me close it