def METHOD_NAME(
    self, x: torch.Tensor, state: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Perform exponential moving average for a single step.
    Args:
        x: MultiHeadDampedEMA input sequences. (B, D, 1)
        state: MultiHeadDampedEMA state. (B, D, num_heads)
    Returns:
        out: MultiHeadDamped output sequences. (B, 1, D)
        new_state: MultiHeadDampedEMA state. (B, D, num_heads)
    """
    damping_factor, prev_timestep_weight = self.get_ema_coefficients()
    new_state = (damping_factor * self.ema_expansion_matrix).squeeze(-1) * x
    if state is not None:
        new_state = new_state + prev_timestep_weight.squeeze(-1) * state
    out = torch.einsum(
        "bdn, dn -> bd", new_state, self.kernel_projection_matrix * self.scaling
    )
    return out.unsqueeze(0), new_state