import torch


def squash(s, dim=-1, eps=1e-8):
	"""
	"Squashing" non-linearity that shrunks short vectors to almost zero 
	length and long vectors to a length slightly below 1
	v_j = ||s_j||^2 / (1 + ||s_j||^2) * s_j / ||s_j||
	
	Args:
		s: 	 Vector before activation
		dim: Dimension along which to calculate the norm
	
	Returns:
		v:   Squashed vector
	"""
	squared_norm = torch.sum(s**2, dim=dim, keepdim=True)
	v = squared_norm / (1 + squared_norm) * \
		s / (torch.sqrt(squared_norm) + eps)

	return v