Log the sizes and shapes of variables, and the total size.

  Args:
    var_list: a list of variables; defaults to trainable_variables
    tag: a string; defaults to "Trainable Variables"
    verbose: bool, if True, log every weight; otherwise, log total size only.
    mesh_to_impl: an optional map from Mesh to MeshImpl

def log_variable_sizes(var_list, tag, verbose=True, mesh_to_impl=None):
  """Log the sizes and shapes of variables, and the total size.

  Args:
    var_list: a list of variables; defaults to trainable_variables
    tag: a string; defaults to "Trainable Variables"
    verbose: bool, if True, log every weight; otherwise, log total size only.
    mesh_to_impl: an optional map from Mesh to MeshImpl
  """
  if not var_list:
    return

  name_to_var = {v.name: v for v in var_list}
  total_size = 0
  total_slice_size = 0
  for v_name in sorted(list(name_to_var)):
    v = name_to_var[v_name]
    v_size = v.shape.size
    if mesh_to_impl is not None:
      slice_size = mesh_to_impl[v.mesh].slice_size(v.shape)
    else:
      slice_size = 0
    total_slice_size += slice_size
    if verbose:
      tf.logging.info(
          "Variable %s size %s slice_size %s %s",
          v.name.ljust(60),
          str(v_size).ljust(12),
          str(slice_size).ljust(12),
          str(v.shape).ljust(60))
      if isinstance(v, StackedVariable):
        for n in v.original_names:
          tf.logging.info("    " + n)
    total_size += v_size
  tf.logging.info("%s count: %s  Total size: %s  Total slice_size: %s",
                  tag.ljust(30), str(len(var_list)).ljust(6),
                  str(total_size).ljust(15),
                  str(total_slice_size).ljust(15))