Implement scaling by AdaBelief about optax HOT 5 CLOSED

I would like to implement this issue!

from optax.

Awesome! Looking forward to your contribution!

Couple of suggestions:

Please move the bias correction code out of scale_by_adam into a standalone fn:

def _bias_correction(moment, decay, count):
  """Perform bias correction. This becomes a no-op as count goes to infinity."""
  bias_correction = 1 - decay**count
  return jax.tree_map(lambda t: t / bias_correction.astype(t.dtype), moment)

Then add both a scale_by_belief in transform.py and a adabelief optimiser in alias.py.

Thanks!

from optax.

Thanks @joaogui1 !

from optax.

We had to briefly revert the change since there was a merging issue,
your PR has now been resubmitted (still with you as author).

from optax.

AdaBelief has been documented #83 thanks everyone 👍

from optax.