41 lines
968 B
Python
41 lines
968 B
Python
from abc import abstractmethod
|
|
from typing import Literal
|
|
|
|
from torch import Tensor
|
|
from torch.nn import Module
|
|
from torch.nn.functional import silu, gelu
|
|
|
|
class GatedUnit(Module):
|
|
def __init__(self, dim: int = -1) -> None:
|
|
super().__init__()
|
|
|
|
self.dim = dim
|
|
|
|
@abstractmethod
|
|
def _activation(self, x: Tensor) -> Tensor:
|
|
...
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
f, g = x.chunk(2, dim=self.dim)
|
|
return self._activation(f) * g
|
|
|
|
class SwiGLU(GatedUnit):
|
|
def __init__(self, dim: int = -1) -> None:
|
|
super().__init__(dim)
|
|
|
|
def _activation(self, x: Tensor) -> Tensor:
|
|
return silu(x)
|
|
|
|
class GeGLU(GatedUnit):
|
|
def __init__(
|
|
self,
|
|
dim: int = -1,
|
|
approximate: Literal["tanh", "none"] = "tanh"
|
|
) -> None:
|
|
super().__init__(dim)
|
|
|
|
self.approximate = approximate
|
|
|
|
def _activation(self, x: Tensor) -> Tensor:
|
|
return gelu(x, self.approximate)
|