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