def _scale_update(update, param):
# Clip norms to minimum value, by default no clipping.
param_norm = numerics.safe_norm(param, min_norm)
update_norm = numerics.safe_norm(update, min_norm)
trust_ratio = trust_coefficient * param_norm / (update_norm + eps)
# If no minimum norm clipping is used
# Set trust_ratio to 1 in case where parameters would never be updated.
zero_norm = jnp.logical_or(param_norm == 0., update_norm == 0.)
safe_trust_ratio = jnp.where(
zero_norm, jnp.array(1.0, dtype=param.dtype), trust_ratio)
return update * safe_trust_ratio
def update_fn(updates, state, params):
if params is None:
raise ValueError(base.NO_PARAMS_MSG)
updates = jax.tree_multimap(
lambda u, p: u * numerics.safe_norm(p, min_scale),
updates, params)
return updates, state
def update_fn(updates, state, params):
updates = jax.tree_multimap(
lambda u, p: u * numerics.safe_norm(p, min_scale), updates, params)
return updates, state