def sum_model_params(
    params: Union[Dict[str, Any], flax.core.FrozenDict]) -> float:
  """Absolute sum of the parameters useful for checking finetuning."""
  all_leaves = jax.tree_util.tree_leaves(params)
  params_sum = []
  for leaf in all_leaves:
    params_sum.append(jnp.sum(jnp.abs(leaf)))
  return np.sum(params_sum, axis=None)
Absolute sum of the parameters useful for checking finetuning.
videogvt/train_lib/train_utils.py