Sun Ruoxi
0c1bb7a415
All checks were successful
Docker Build and Push / docker (push) Successful in 1m5s
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
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
|