import torch


def neigh_square(position_diff: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
    """
    Square-shaped neighborhood scaling function based on center-wise diff `position_diff` and radius `sigma`.

    Parameters
    ----------
    position_diff : torch.Tensor
        The positional difference around some center.
    sigma : torch.Tensor
        The scaling radius.
    """
    v = (position_diff).abs().max(-1)[0]
    # v = pnorm(position_diff, p=2)
    return torch.exp(torch.neg(v.sqrt() / sigma))