import torch


def hamming_distance(logits, targets, cost_scale=1.0):
    """
    Computes Hamming distance (https://en.wikipedia.org/wiki/Hamming_distance), which is
    defined as the number of positions where two sequences of equal length differ. We apply
    Hamming distance locally, incrementing non-gold token scores by `cost_scale`.

    ```
    Example:
    Given targets = [0, 1] and cost_scale = 1.0, we have the following:
    logits (before) = [[-1.0, 1.0, 2.0], [-2.0, -1.0, 1.0]]
    logits (after) = [[-1.0, 2.0, 3.0], [-1.0, -1.0, 2.0]]
    ```
    """

    hamming_cost = cost_scale * torch.ones_like(logits)  # B x T x V
    gold_cost = torch.zeros_like(targets).to(logits.dtype).unsqueeze(2)  # B x T x 1
    hamming_cost.scatter_(2, targets.unsqueeze(2), gold_cost)

    return hamming_cost