import torch


def _tanh_gate(x: torch.Tensor, alpha, beta, inverse=False):
    """ Invertible point-wise nonlinearity implemented by mixture of linear and tanh

        Arguments:
            x:      input tensor
            alpha:  bandwidth factor for tanh (0 < alpha)
            beta:   mixing coefficient interpolating between linear and tanh (0 < beta < 1)

        Returns:
            y:      transformed data
                    torch.Tensor
            dlogp:  log of diagonal jacobian
                    torch.Tensor
    """
    if not inverse:
        dlogp = torch.log(
            beta + (1.0 - beta) * alpha * (1.0 - torch.tanh(alpha * x).pow(2))
        )
        y = beta * x + (1.0 - beta) * torch.tanh(alpha * x)
        return y, dlogp
    else:
        raise NotImplementedError()