import torch


def swish(x: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:
  """Swish activation function
  from arXiv:1710.05941
  
  Args:
    x (torch.Tensor): input tensors can be either 2D or 4D.
    beta (torch.Tensor): swish-beta can be a constant or a trainable params.
  
  Returns:
    torch.Tensor: output tensors
  """
  if len(x.size()) == 2:
    return x * torch.sigmoid(beta[None, :] * x)
  else:
    return x * torch.sigmoid(beta[None, :, None, None] * x)