def _convert_to_parameter_stats(global_stats, local_stat, convert_statistics=True):
    """Creates parameter stats from sharded stats."""
    index_start = int(local_stat.index_start)
    index_end = int(len(local_stat.sizes)) + index_start
    statistics = global_stats.statistics[index_start:index_end, :, :]
    preconditioners = global_stats.preconditioners[index_start:index_end, :, :]
    new_statistics = []
    new_preconditioners = []
    for i, size in enumerate(local_stat.sizes):
        new_statistics.append(statistics[i][:size, :size])
        new_preconditioners.append(preconditioners[i][:size, :size])
    if not convert_statistics:
        new_statistics = None
    return ParameterStats(
        local_stat.diagonal_statistics,
        new_statistics,
        new_preconditioners,
        local_stat.diagonal_momentum,
        local_stat.momentum,
        local_stat.training_metrics,
    )
Creates parameter stats from sharded stats.
tools/train/scalable_shampoo/distributed_shampoo.py