import torch
import torch.nn as nn


def get_activation_class(activation: str, dim=None):
    """ Returns the activation function corresponding to `activation` """

    if activation == "relu":
        return nn.ReLU()
    elif activation == "gelu":
        return nn.GELU()
    elif activation == "glu":
        assert dim is not None
        return nn.GLU(dim=dim)
    elif activation == "swish":
        return Swish()
    elif activation == "none":
        return nn.Identity()
    else:
        raise RuntimeError("activation function {} not supported".format(activation))

def swish(x: torch.Tensor) -> torch.Tensor:
    return x * torch.sigmoid(x)


class Swish(nn.Module):

    def __init__(self):
        super(Swish, self).__init__()

    def forward(self, x):
        return x * x.sigmoid()
