How to resolve `if/else' with multiple statement for jax.grad()??

[yyang97]

import jax
import jax.numpy as jnp
import jax.scipy as jsp
import jax.random
import jax.test_util
# from jaxopt import Bisection
from jax import lax, jit
import jaxopt


from jax.config import config
config.update("jax_enable_x64", True)


import scipy
import numpy as np

def gamma(z):
    return jnp.exp(jax.scipy.special.gammaln(z))


def const_Tc(Ts, T0, pars, k):

    A = pars[0]
    b = pars[1]
    C = pars[2]
    n = pars[3]
    s0 = pars[4]
    mu = 1
    
    As = A*k*Ts
    Cs = C*k*Ts
    
    intfac = 1/mu * jnp.power(Cs,n)*Ts/(n+1)*jnp.power(T0/Ts-s0,n+1)
    bnfr = (b+1)/(n+1)
    
    res = jnp.zeros(4)
    res = res.at[0].set(1/mu * jnp.power(As*jnp.maximum(T0/Ts - s0, 0.0),b))
    res = res.at[1].set(1/mu * jnp.power(Cs*jnp.maximum(T0/Ts - s0, 0.0),n))
    res = res.at[2].set(jnp.exp(-res[1] * T0 + intfac) *jnp.power(As,b) / jnp.power(Cs,n*bnfr) * jnp.power(mu/Ts * (n+1),(b-n)/(n+1)) * jsp.special.gammainc(bnfr,intfac)*gamma(bnfr))
    res = res.at[3].set(jnp.exp(-res[1] * T0))
    Tc = -1/res[1] * jnp.log((res[0]/res[1]*res[3] + res[2]) / (1 + res[0]/res[1]))
    
    return Tc   

My code works if

# it works
T0 = 0.011584800741427247
k = 388440
Ts = 0.01642112
pars = jnp.array([4.59185910e-04, 2.58638497e+01, 2.64082837e-10, 3.17472539e-01,
        5.31952602e-01])
const_Tc(Ts, T0, pars, k)
# Array(1510.59335656, dtype=float64)
jax.grad(const_Tc,2)(Ts, T0, pars, k)
# Array([-6.68528443e+06,  8.03246190e+01, -1.67329068e+12,  2.09416501e+04,
        2.02368051e+04], dtype=float64)

It returns nan if

# it does not works
T0 = 0.011584800741427247
k = 388440
Ts = 0.02353615
pars = jnp.array([3.25051444e-04, 2.53327831e+01, 3.92844549e-10, 3.94831369e-01,
        5.42143459e-01])
# it returns nan
const_Tc(Ts, T0, pars, k)
# it returns nan
jax.grad(const_Tc,2)(Ts, T0, pars, k)


s0 = pars[4]
n = pars[3]
# nan
jnp.power(T0/Ts-s0,n+1)

The reason is that my T0/Ts-s0 < 0 and it returns nan.

I hope my const_Tc() return 0 if T0/Ts-s0 < 0. But how can I write it in jax?

# def const_Tc(Ts, T0, pars, k):

#     A = pars[0]
#     b = pars[1]
#     C = pars[2]
#     n = pars[3]
#     s0 = pars[4]
#     mu = 1
    
#     As = A*k*Ts
#     Cs = C*k*Ts
    
#     if T0/Ts-s0>0:
#         intfac = 1/mu * jnp.power(Cs,n)*Ts/(n+1)*jnp.power(T0/Ts-s0,n+1)
#         bnfr = (b+1)/(n+1)

#         res = jnp.zeros(4)
#         res = res.at[0].set(1/mu * jnp.power(As*jnp.maximum(T0/Ts - s0, 0.0),b))
#         res = res.at[1].set(1/mu * jnp.power(Cs*jnp.maximum(T0/Ts - s0, 0.0),n))
#         res = res.at[2].set(jnp.exp(-res[1] * T0 + intfac) *jnp.power(As,b) / jnp.power(Cs,n*bnfr) * jnp.power(mu/Ts * (n+1),(b-n)/(n+1)) * jsp.special.gammainc(bnfr,intfac)*gamma(bnfr))
#         res = res.at[3].set(jnp.exp(-res[1] * T0))
#         Tc = -1/res[1] * jnp.log((res[0]/res[1]*res[3] + res[2]) / (1 + res[0]/res[1]))

#         return Tc   
#     elif:
#         return 0 

This is what I really want to have. But how can I write it use lax.cond() or jnp.select()?

[jakevdp]

See JAX Sharp Bits: Control Flow.

There are a few ways you could express this (including lax.cond, lax.select, etc.). In your case, it may be easiest to use jnp.where, something like this:

def const_Tc(Ts, T0, pars, k):

    A = pars[0]
    b = pars[1]
    C = pars[2]
    n = pars[3]
    s0 = pars[4]
    mu = 1
    
    As = A*k*Ts
    Cs = C*k*Ts
    
    intfac = 1/mu * jnp.power(Cs,n)*Ts/(n+1)*jnp.power(T0/Ts-s0,n+1)
    bnfr = (b+1)/(n+1)

    res = jnp.zeros(4)
    res = res.at[0].set(1/mu * jnp.power(As*jnp.maximum(T0/Ts - s0, 0.0),b))
    res = res.at[1].set(1/mu * jnp.power(Cs*jnp.maximum(T0/Ts - s0, 0.0),n))
    res = res.at[2].set(jnp.exp(-res[1] * T0 + intfac) *jnp.power(As,b) / jnp.power(Cs,n*bnfr) * jnp.power(mu/Ts * (n+1),(b-n)/(n+1)) * jsp.special.gammainc(bnfr,intfac)*gamma(bnfr))
    res = res.at[3].set(jnp.exp(-res[1] * T0))
    Tc = -1/res[1] * jnp.log((res[0]/res[1]*res[3] + res[2]) / (1 + res[0]/res[1]))

    return jnp.where(T0/Ts-s0>0, Tc, 0)

[yyang97]

Thank you very much for your reply. Yeah it does gives 0 for this function.

My apology is that what I want is it can works for jax.grad().

# it does not works
T0 = 0.011584800741427247
k = 388440
Ts = 0.02353615
pars = jnp.array([3.25051444e-04, 2.53327831e+01, 3.92844549e-10, 3.94831369e-01,
        5.42143459e-01])
# it returns 0
const_Tc(Ts, T0, pars, k)
# it returns nan
jax.grad(const_Tc,2)(Ts, T0, pars, k)

Using the new function, jax.grad() still returns nan...

[jakevdp]

The graient issue is discussed here: https://jax.readthedocs.io/en/latest/faq.html#gradients-contain-nan-where-using-where

You can fix it by modifying the problematic value at the start of the computation. This will be easiest to do if you know which computation the NaN value is coming from – do you know which it is?