Iurii Kemaev

@joaogui1 it should be fixed now, could you update the PR?

>Should I merge main? Yes 👍

@Zach-ER thanks a lot for your great comments! Since you are an experienced user of this optimiser, I was wondering if the current version looks good to you? :)

Thanks for the question! @mtthss could you comment on this, since you're the one who ported Adafactor from Flax?