def _add_error_into_local_stats(local_stats, errors, inverse_failure_threshold):
    """Adds errors back into local statistics."""
    new_local_stats = []
    for local_stat in local_stats:
        if local_stat.sizes:
            index_start = int(local_stat.index_start)
            index_end = int(len(local_stat.sizes)) + index_start
            per_stat_error = errors[index_start:index_end]
        else:
            per_stat_error = jnp.array(0, jnp.float32)
        if local_stat.sizes:
            per_stat_error = jnp.where(
                jnp.logical_and(
                    per_stat_error > 0.0, per_stat_error != inverse_failure_threshold
                ),
                per_stat_error,
                local_stat.training_metrics.inverse_pth_root_errors,
            )
        new_local_stats.append(
            LocalShardedParameterStats(
                local_stat.diagonal_statistics,
                local_stat.diagonal_momentum,
                local_stat.momentum,
                TrainingMetrics(per_stat_error),
                local_stat.index_start,
                local_stat.sizes,
            )
        )
    return new_local_stats
Adds errors back into local statistics.
tools/train/scalable_shampoo/distributed_shampoo.py