import torch


def reg_psd(X: torch.Tensor, reg: float = 1e-3) -> torch.Tensor:
    """Regularizes a batch of PSD matrix trajectories."""
    x_dim = X.shape[-1]
    return X + reg * torch.eye(x_dim, device=X.device)