Solution for ad of mcmc

mcad.py

@tf.custom_gradient
def idn(x):
    def grad(*dx):
        return tf.zeros_like(dx[0])
    return x, grad

def probexp(beta, s):
    return tf.exp(-beta*s)

def obss(beta, s):
    return s**2*beta

def updateexp(s):
    step = 5.8
    while True:
        r = tf.random.uniform(shape=[1], maxval=step, minval=-step)
        if s+r>0:
            break
    return s+r
    
def evalobs(prob, obs, update, beta, s, times):
    l0, l1, l2, l3 = 0, 0, 0, 0
    reject = 0
    for _ in range(times):
        snew = update(s)
        ratio = prob(beta, snew)/prob(beta, s)
        if ratio>1:
            s = snew
        elif tf.random.uniform(maxval=1, minval=0, shape=[1])<ratio:
            s = snew
        else:
            reject += 1
        l0 += obs(beta, s)

        l1 += idn(obs(beta, s))*tf.log(prob(beta, s))
        l2 += idn(obs(beta, s))
        l3 += tf.log(prob(beta, s))
    deltal = l1/times-l2*l3/times**2
    l = l0/times+deltal-idn(deltal)
    print(reject)
    return l