Generalizes safe_int32_increment to safe_increment

[copybara-service]

Generalizes safe_int32_increment to safe_increment

[carlosgmartin]

@vroulet I see that optax.safe_increment rejects unsigned integers, and its code contains the following rationale as a comment:

The comparison count < max_value appears to convert its arguments into signed integers so we get overflow errors with unsigned integers.

It looks like performing the comparison against an array containing max_value with the dtype of count solves this problem:

>>> from jax import numpy as jnp
>>> count = jnp.array(0, jnp.uint32)
>>> max_val = jnp.iinfo(count.dtype).max
>>> count < max_val
OverflowError: Python int 4294967295 too large to convert to int32
>>> count < jnp.array(max_val, count.dtype)
Array(True, dtype=bool)

If that's a valid solution, I can submit a PR for this.

[vroulet]

Ah I should have thought about this.
You can send a PR to fix this (add corresponding tests). Thanks for catching this!

[carlosgmartin]

@vroulet Done: #1092.