JITted Adam results in NaN when setting decay to integer 0 about optax

[mtthss]

Interesting, thanks for reaching out!
Let me see what the timeline would be to fix this in core jax itself

Just a note:
if you use Python's type checking or if you are using an IDE like PyCharm that supports Python type checking,
passing an int to adam would actually raise an error without even having to run the code

Since we explicitly type both the b1 and b2 arguments as float.

def scale_by_adam(b1: float = 0.9,
                  b2: float = 0.999,
                  eps: float = 1e-8,
                  eps_root: float = 0.0) -> GradientTransformation:

[mtthss]

A fix should have gone in in JAX

[dlsmith]

Thanks for following up. Yea, I absolutely agree on the use of type checking, though it would be unfortunate if correctness relied on opting in to third party tooling, and in a case where assuming promotion semantics fails in a fairly subtle way. In any case, great that the fix came in so fast!

[mtthss]

Yep indeed they were quite quick!
Closing the issue since this should now be fixed