import torch


def squash(x, dim):
    """ Non-linear activation, that squashes all vectors to have norm < 1 """
    norm_sq = torch.sum(x ** 2, dim, keepdim=True)
    norm = torch.sqrt(norm_sq)
    return (norm_sq / (1.0 + norm_sq)) * (x / norm)