import torch


def _leaky_relu_gate(x: torch.Tensor, a, b, inverse=False):
    """ Invertible point-wise nonlinearity implemented by leaky relu

        Arguments:
            x: input tensor
            a: positive scaling factor (0 < a)
            b: negative scaling factor (0 < b)

        Returns:
            y:      transformed data
                    torch.Tensor
            dlogp:  log of diagonal jacobian
                    torch.Tensor
    """
    cond = x >= 0
    if not inverse:
        y = cond * x * a + ~cond * x * b
        dlogp = torch.log(cond * a + ~cond * b)
        return y, dlogp
    else:
        y = cond * x / a + ~cond * x / b
        dlogp = torch.log(cond / a + ~cond / b)
        return y, dlogp