from typing import Dict
from typing import Any


def clear_per_step_extra_state(extra_state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Clear values in extra_state that are technically only true for a specific
    step (ex: the eval tune loss calculated after 5 train steps is no longer
    accurate after 7 train steps, but might not get updated since we might not
    be doing eval after every step).
    """
    extra_state["tune_eval"]["loss"] = None
    extra_state["tune_eval"]["perplexity"] = None
    extra_state["tune_bleu"]["current"] = None
    return extra_state