Change the grads argument in optax.apply_if_finite.update to updates, so that it is consistent with other update functions (e.g., optax.clip_by_global_norm.update).
Change the grads argument in optax.apply_if_finite.update to updates, so that it is consistent with other update functions (e.g., optax.clip_by_global_norm.update).

PiperOrigin-RevId: 351379232
Change-Id: Icf641ef6a6851d4020e44c24f0d815a44a5bd1ab

deepmind/optax,deepmind/optax