30 lines
585 B
Python
30 lines
585 B
Python
|
import torch
|
||
|
|
|
||
|
|
__all__ = ["gelu_tanh_and_mul"]
|
||
|
|
|
||
|
|
|
||
|
|
def gelu_tanh_and_mul(
|
||
|
|
input: "torch.Tensor",
|
||
|
|
output: "torch.Tensor" = None
|
||
|
|
):
|
||
|
|
assert isinstance(input, torch.Tensor)
|
||
|
|
|
||
|
|
if output is None:
|
||
|
|
output_shape = list(input.shape)
|
||
|
|
output_shape[-1] = output_shape[-1] // 2
|
||
|
|
output = input.new_empty(output_shape)
|
||
|
|
|
||
|
|
hidden_size = input.shape[-1] // 2
|
||
|
|
|
||
|
|
x = input[..., :hidden_size]
|
||
|
|
gate = input[..., hidden_size:]
|
||
|
|
|
||
|
|
result = x * torch.nn.functional.gelu(
|
||
|
|
gate,
|
||
|
|
approximate="tanh"
|
||
|
|
)
|
||
|
|
|
||
|
|
output.copy_(result)
|
||
|
|
|
||
|
|
return output
|