import torch
from typing import Optional


def squash(inputs: torch.Tensor, dim: Optional[int] = -1) -> torch.Tensor:
    """
    Apply `squash` non-linearity to inputs
    
    Args:
        inputs (T): input tensor which is to be applied squashing
        dim (int, optional): dimension to be applied squashing. Defaults to -1
    
    Returns:
        T: squashed tensor
    """
    # Calculate squared norm of inputs
    squared_norm = torch.sum(torch.pow(inputs, 2), dim=dim, keepdim=True)

    # Squash non-linearity to inputs
    return (squared_norm / (1 + squared_norm)) * (inputs / (torch.sqrt(squared_norm) + 1e-8))