import torch


def squash(tensor):
    """
    Squash function, defined in [1]. Works as a non-linearity for CapsNets.
    Input tensor will be of format (bs, units, C, H, W) or (bs, units, C)
    Norm should be computed on the axis representing the number of units.

    Parameters
    ----------
        tensor :    torch Variable containing n-dimensional tensor

    Returns
    -------
        (||tensor||^2 / (1+ ||tensor||^2)) * tensor/||tensor||
    """
    norm = torch.norm(tensor, p=2, dim=1, keepdim=True)
    sq_norm = norm ** 2  # Avoid computing square twice

    return tensor.div(norm) * sq_norm / (1 + sq_norm)