import torch


def l0_ohe(input: torch.Tensor,
           target: torch.Tensor,
           reduction: str = 'none') -> torch.Tensor:
    """
    Computes the L0 loss for a one-hot encoding representation.

    Parameters
    ----------
    input
        Input tensor.
    target
        Target tensor
    reduction
        Specifies the reduction to apply to the output: `none` | `mean` | `sum`.

    Returns
    -------
        L0 loss.
    """
    # Order matters as the gradient of zeros will still flow if reversed order. Maybe consider clipping a bit higher?
    eps = 1e-7 / input.shape[1]
    loss = torch.maximum(target - input, eps + torch.zeros_like(input)).sum(dim=1)

    if reduction == 'none':
        return loss

    if reduction == 'mean':
        return torch.mean(loss)

    if reduction == 'sum':
        return torch.sum(loss)

    raise ValueError(f"Reduction {reduction} not implemented.")