import torch


def invSquare(input, axis=1):
    """ 
    Apply inverse square normalization on input at certain axis.
    
    Parammeters:
    ----------
    input: Tensor (N*L or rank>2)
    axis: the axis to apply softmax
    
    Returns: Tensor with softmax applied on that dimension.
    """
    
    input_size = input.size()
    
    trans_input = input.transpose(axis, len(input_size)-1)
    trans_size = trans_input.size()

    input_2d = trans_input.contiguous().view(-1, trans_size[-1])
    
    square_2d = input_2d**(-2)
    sum_square_2d = torch.sum(square_2d, 1, keepdim=True)
    square_norm_2d = square_2d/sum_square_2d

    square_norm_nd = square_norm_2d.view(*trans_size)
    return square_norm_nd.transpose(axis, len(input_size)-1)