"""IRMv1 penalty (Arjovsky et al., 2019) -- DG baseline.

Penalizes the sensitivity of the per-environment risk to a dummy scalar classifier
of value 1: ``(d/dw CE(w * logits, y))^2`` at ``w = 1``. When the representation is
truly invariant the optimal classifier is simultaneously optimal on every
environment, so this gradient vanishes. Rarely beats tuned ERM (Gulrajani &
Lopez-Paz); included as a required baseline (ARCV_METHOD_DESIGN.md §3, row 8).
"""
import torch
import torch.nn.functional as F


def irm_penalty(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    scale = torch.ones(1, device=logits.device, requires_grad=True)
    loss = F.cross_entropy(logits * scale, labels)
    grad, = torch.autograd.grad(loss, scale, create_graph=True)
    return (grad ** 2).sum()