r""" GELU activation function with tanh approximation support with `approximate="tanh"`. Parameters: dim_in (`int`): The number of channels in the input. dim_out (`int`): The number of channels in the output. approximate (`str`, *optional*, defaults to `"none"`): If
| 63 | |
| 64 | |
| 65 | class GELU(nn.Module): |
| 66 | r""" |
| 67 | GELU activation function with tanh approximation support with `approximate="tanh"`. |
| 68 | |
| 69 | Parameters: |
| 70 | dim_in (`int`): The number of channels in the input. |
| 71 | dim_out (`int`): The number of channels in the output. |
| 72 | approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation. |
| 73 | bias (`bool`, defaults to True): Whether to use a bias in the linear layer. |
| 74 | """ |
| 75 | |
| 76 | def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True): |
| 77 | super().__init__() |
| 78 | self.proj = nn.Linear(dim_in, dim_out, bias=bias) |
| 79 | self.approximate = approximate |
| 80 | |
| 81 | def gelu(self, gate: torch.Tensor) -> torch.Tensor: |
| 82 | if gate.device.type == "mps" and is_torch_version("<", "2.0.0"): |
| 83 | # fp16 gelu not supported on mps before torch 2.0 |
| 84 | return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(dtype=gate.dtype) |
| 85 | return F.gelu(gate, approximate=self.approximate) |
| 86 | |
| 87 | def forward(self, hidden_states): |
| 88 | hidden_states = self.proj(hidden_states) |
| 89 | hidden_states = self.gelu(hidden_states) |
| 90 | return hidden_states |
| 91 | |
| 92 | |
| 93 | class GEGLU(nn.Module): |