Diff Coverage

                Both modes support distinct update and buffer coefficients.
        """
        advanced = _resolve_muon_advanced_options(advanced_options)
        if momentum_mode not in ("legacy", "vanilla"):
            raise ValueError(
                f"momentum_mode must be 'legacy' or 'vanilla', got {momentum_mode!r}"
            )
        if ns_variant not in ("legacy", "asym5", "custom"):
            raise ValueError(

            # m_for_update = (1 - momentum1) * grad + momentum1 * m_old
            # m_new = (1 - momentum2) * grad + momentum2 * m_old
            # Fold the Nesterov interpolation into one lerp with momentum1 squared.
            # pylint: disable=protected-access
            update_momentum = momentum1 ** 2 if group["nesterov"] else momentum1
            local_us = torch._foreach_lerp(grads, bufs, weight=update_momentum)
            torch._foreach_lerp_(bufs, grads, weight=1 - momentum2)
        elif self._update_muon_momentum_mode == "legacy":
            # Match muon_update_core():
            # m_for_update = grad + momentum1 * m_old
            # m_new = grad + momentum2 * m_old

            torch._foreach_mul_(bufs, momentum2)
            torch._foreach_add_(bufs, grads)

            if group["nesterov"]:
                torch._foreach_mul_(local_us, momentum1)
                torch._foreach_add_(local_us, grads)
        else:
            raise ValueError(
                "Muon._update_muon_momentum_mode must be 'vanilla' or 'legacy', "
                f"got {self._update_muon_momentum_mode!r}"
            )