def jacobian_det(self, y):
     Km1 = y.shape[0]
     k = tt.arange(Km1)[(slice(None),) + (None,) * (y.ndim - 1)]
     eq_share = -tt.log(Km1 - k)  # logit(1./(Km1 + 1 - k))
     yl = y + eq_share
     yu = tt.concatenate([tt.ones(y[:1].shape), 1 - inverse_logit(yl)])
     S = tt.extra_ops.cumprod(yu, 0)
     return tt.sum(tt.log(S[:-1]) - tt.log1p(tt.exp(yl)) - tt.log1p(tt.exp(-yl)), 0)

 def jacobian_det(self, y_):
     y = y_.T
     Km1 = y.shape[0]
     k = tt.arange(Km1)[(slice(None), ) + (None, ) * (y.ndim - 1)]
     eq_share = logit(1. / (Km1 + 1 - k).astype(str(y_.dtype)))
     yl = y + eq_share
     yu = tt.concatenate([tt.ones(y[:1].shape), 1 - invlogit(yl, self.eps)])
     S = tt.extra_ops.cumprod(yu, 0)
     return tt.sum(tt.log(S[:-1]) - tt.log1p(tt.exp(yl)) - tt.log1p(tt.exp(-yl)), 0).T

    def rates(self, x):
        dtype = theano.config.floatX
        sigma = tt.cast(0.05, dtype=dtype)
        tau_ref = tt.cast(self.tau_ref, dtype=dtype)
        tau_rc = tt.cast(self.tau_rc, dtype=dtype)

        j = self.gain * x + self.bias - 1
        j = sigma * tt.log1p(tt.exp(j / sigma))
        v = 1. / (tau_ref + tau_rc * tt.log1p(1. / j))
        return tt.switch(j > 0, v, 0.0) / self.max_rates

def log_i0(x):
    """
    Calculates the logarithm of the 0 order modified Bessel function of the first kind""
    """
    return tt.switch(tt.lt(x, 5), tt.log1p(x**2. / 4. + x**4. / 64. + x**6. / 2304.
                                           + x**8. / 147456. + x**10. / 14745600.
                                           + x**12. / 2123366400.),
                                  x - 0.5 * tt.log(2. * np.pi * x) + tt.log1p(1. / (8. * x)
                                  + 9. / (128. * x**2.) + 225. / (3072. * x**3.)
                                  + 11025. / (98304. * x**4.)))

def safe_logaddexp(a, b):
    """Symbolic log(exp(a) + exp(b)). The edge case where `a` - `b` is undefined is handled by
    setting the difference to 0. This occurs if both `a` and `b` are +inf or -inf.

    Returns:
        symbolic log(exp(a) + exp(b))
    """
    diff = b - a
    safe_diff = tt.switch(tt.isnan(diff), 0, diff)
    return tt.switch(safe_diff >= 0,
                     b + tt.log1p(tt.exp(-safe_diff)),
                     a + tt.log1p(tt.exp(safe_diff)))

    def logp(self, value):
        psi = self.psi
        p = self.p
        n = self.n

        logp_val = tt.switch(tt.gt(value, 0),
                 tt.log(psi) + self.bin.logp(value),
                 logaddexp(tt.log1p(-psi), tt.log(psi) + n * tt.log1p(-p)))

        return bound(logp_val,
            0 <= value, value <= n,
            0 <= psi, psi <= 1,
            0 <= p, p <= 1)

def nlif(x):
    dtype = theano.config.floatX
    sigma = tt.cast(0.05, dtype=dtype)
    tau_ref = tt.cast(0.002, dtype=dtype)
    tau_rc = tt.cast(0.02, dtype=dtype)
    alpha = tt.cast(1, dtype=dtype)
    beta = tt.cast(1, dtype=dtype)
    amp = tt.cast(1. / 65, dtype=dtype)

    j = alpha * x + beta - 1
    j = sigma * tt.log1p(tt.exp(j / sigma))
    v = amp / (tau_ref + tau_rc * tt.log1p(1. / j))
    return tt.switch(j > 0, v, 0.0)

def nlif(x):
    dtype = theano.config.floatX
    sigma = tt.cast(0.05, dtype=dtype)
    tau_ref = tt.cast(0.002, dtype=dtype)
    tau_rc = tt.cast(0.02, dtype=dtype)
    alpha = tt.cast(1, dtype=dtype)
    beta = tt.cast(1, dtype=dtype)  # so that f(0) = firing threshold
    amp = tt.cast(1. / 63.04, dtype=dtype)  # so that f(1) = 1

    j = alpha * x + beta - 1
    j = sigma * tt.log1p(tt.exp(j / sigma))
    v = amp / (tau_ref + tau_rc * tt.log1p(1. / j))
    return tt.switch(j > 0, v, 0.0)

 def kl_div(self, x, y):
     """
     Compute sum of D(x_i || y_i) for each corresponding element
     along the 3rd dimension (the embedding dimension)
     of x and y
     This function takes care to not compute logarithms that are close
     to 0, since NaN's could result for log(sigmoid(x)) if x is negative.
     It simply uses that log(sigmoid(x)) = - log(1 + e^-x)
     """
     sig_x = T.nnet.sigmoid(x)
     exp_x = T.exp(x)
     exp_y = T.exp(y)
     exp_neg_x = T.exp(-x)
     exp_neg_y = T.exp(-y)
     return (sig_x * (T.log1p(exp_neg_y) - T.log1p(exp_neg_x)) + (1 - sig_x) * (T.log1p(exp_y) - T.log1p(exp_x))).mean()

    def __call__(self, layer, spec, shape, name=None, **tags):
        # case when user uses default init specs
        assert tags.get('variational',False) == True, "Please declare param as variational to avoid confusion"
        
        if not isinstance(spec, dict):
            initial_rho = np.log(np.expm1(self.prior_std))   #std to rho
            assert np.isfinite(initial_rho),"too small std to initialize correctly. Please pass explicit"\
                                            " initializer (dict with {'mu':mu_init, 'rho':rho_init})."
            spec = {'mu': spec,'rho':init.Constant(initial_rho)}
            

        mu_spec,rho_spec = spec['mu'],spec['rho']
        
        rho = layer.add_param(rho_spec, shape,name=(name or 'unk')+'.rho', **tags)
        mean = layer.add_param(mu_spec, shape,name=(name or 'unk')+'.mu', **tags)

        #Reparameterization trick
        e = self.srng.normal(shape, std=1)  
        W = mean + T.log1p(T.exp(rho)) * e 

        #KL divergence KL(q,p) = E_(w~q(w|x)) [log q(w|x) - log P(w)] aka variational cost
        q_p = T.sum(self.log_posterior_approx(W, mean, rho) - self.log_prior(W))
            
        #accumulate variational cost
        layer._bbwrap_var_cost += q_p
        return W

 def log_posterior_approx(self,weights, mean, rho):
     """
     Logarithm of ELBO on posterior probabilities:
     log q(weights|learned mu and rho) aka log q(theta|x)
     """
     std = T.log1p(T.exp(rho))  #rho to std
     return self.log_normal(weights, mean, std)

    def setup(self, bottom, top):
        from caffe_helper.theano_util import init_theano
        init_theano()

        import theano as tn
        import theano.tensor as T
        assert len(bottom) == 2
        assert len(top) == 1
        s_y = T.matrix('y')  # y in [-inf, inf]
        s_t = T.matrix('t')  # t in {-1, 0, 1} where 0 is ignored
        s_dloss = T.scalar('dloss')
        # Forward
        # s_loss = T.mean(abs(s_t) * T.log1p(T.exp(-s_y * s_t)))  # unstable
        s_loss = -T.sum(
            abs(s_t) * (
                s_y * ((s_t >= 0) - (s_y >= 0)) - T.log1p(T.exp(-abs(s_y)))))\
            / T.maximum(T.sum(abs(s_t)), 1)
        # Backward
        s_p = 1 / (1 + T.exp(-s_y))
        s_dy = s_dloss * abs(s_t) * (s_p - (s_t >= 0)) / \
            T.maximum(T.sum(abs(s_t)), 1)

        def _o(s):
            return tn.Out(s, borrow=True)
        self.tn_forward = tn.function([s_y, s_t], s_loss)
        self.tn_backward = tn.function([s_y, s_t, s_dloss], _o(s_dy))

def normal_lccdf(mu, sigma, x):
    z = (x - mu) / sigma
    return tt.switch(
        tt.gt(z, 1.0),
        tt.log(tt.erfcx(z / tt.sqrt(2.)) / 2.) - tt.sqr(z) / 2.,
        tt.log1p(-tt.erfc(-z / tt.sqrt(2.)) / 2.)
    )

def normal_lcdf(mu, sigma, x):
    """Compute the log of the cumulative density function of the normal."""
    z = (x - mu) / sigma
    return tt.switch(
        tt.lt(z, -1.0),
        tt.log(tt.erfcx(-z / tt.sqrt(2.)) / 2.) - tt.sqr(z) / 2.,
        tt.log1p(-tt.erfc(z / tt.sqrt(2.)) / 2.)
    )

    def logp(self, value):
        quaddist, logdet, ok = self._quaddist(value)
        k = value.shape[-1].astype(theano.config.floatX)

        norm = (gammaln((self.nu + k) / 2.)
                - gammaln(self.nu / 2.)
                - 0.5 * k * floatX(np.log(self.nu * np.pi)))
        inner = - (self.nu + k) / 2. * tt.log1p(quaddist / self.nu)
        return bound(norm + inner - logdet, ok)

def log_add(lna, lnb):
    """
    Compute the ln(a+b) given {lna,lnb}
    :param
    :return: ln(a+b)
    """
    max_ = tensor.maximum(lna, lnb)
    result = (max_ + tensor.log1p(tensor.exp(lna + lnb - 2 * max_)))   #log1p(x) = log(1+x)
    return tensor.switch(tensor.isnan(result), max_, result)

    def logp(self, value):
        nu = self.nu
        mu = self.mu
        lam = self.lam
        sd = self.sd

        return bound(gammaln((nu + 1.0) / 2.0)
                     + .5 * T.log(lam / (nu * np.pi))
                     - gammaln(nu / 2.0)
                     - (nu + 1.0) / 2.0 * T.log1p(lam * (value - mu)**2 / nu),
                     lam > 0, nu > 0, sd > 0)

def lif(x):
    dtype = theano.config.floatX
    tau_ref = tt.cast(0.002, dtype=dtype)
    tau_rc = tt.cast(0.02, dtype=dtype)
    alpha = tt.cast(1, dtype=dtype)
    beta = tt.cast(1, dtype=dtype)
    amp = tt.cast(1. / 65, dtype=dtype)

    j = alpha * x + beta - 1
    v = amp / (tau_ref + tau_rc * tt.log1p(1. / j))
    return tt.switch(j > 0, v, 0.0)

def _log_add_3(log_a, log_b, log_c):
    """Theano expression for log(a+b+c) given log(a), log(b), log(c)."""
    smaller = T.minimum(log_a, log_b)
    larger = T.maximum(log_a, log_b)
    largest = T.maximum(larger, log_c)
    larger = T.minimum(larger, log_c)

    return largest + T.log1p(
            T.exp(smaller - largest) + 
            T.exp(larger - largest)
            )

def log1mexp(x):
    """Return log(1 - exp(-x)).

    This function is numerically more stable than the naive approch.

    For details, see
    https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf
    """
    return tt.switch(
        tt.lt(x, 0.683),
        tt.log(-tt.expm1(-x)),
        tt.log1p(-tt.exp(-x)))