Add optional weight decay to AdaBelief optimizer

**Feature request**: Add an optional `weight_decay` argument, as seen for other optimizers, to [optax.adabelief](https://optax.readthedocs.io/en/latest/api/optimizers.html#optax.adabelief).