def _convert_from_parameter_stats(parameter_stats, local_stats):
    """Creates sharded stats from paramter stats."""
    return LocalShardedParameterStats(
        parameter_stats.diagonal_statistics,
        parameter_stats.diagonal_momentum,
        parameter_stats.momentum,
        parameter_stats.training_metrics,
        local_stats.index_start,
        local_stats.sizes,
    )
Creates sharded stats from paramter stats.
tools/train/scalable_shampoo/distributed_shampoo.py