from typing import Sequence
from typing import Callable
import torch


def zero(units: Sequence[int]) -> Callable[[torch.Tensor], torch.Tensor]:
    """Zero the given units.

    Args:
        units (Sequence[int]): The units to zero.

    Returns:
        Callable[[torch.Tensor], torch.Tensor]: Function that takes layer
            features and zeros the given units, returning the result.

    """

    def fn(features: torch.Tensor) -> torch.Tensor:
        if features.dim() != 4:
            raise ValueError(f'expected 4D features, got {features.dim()}')
        # Make sure we don't break autograd by editing values in place.
        # Just use a mask. Fauci said it first.
        shape = (*features.shape[:2], 1, 1)
        mask = features.new_ones(*shape)
        mask[:, units] = 0
        return features * mask

    return fn