<jupyter_start><jupyter_text>Problematic gradient breaks the sampler<jupyter_code>f = the_model.logp_dlogp_function()
f.set_extra_values({})
f(f.dict_to_array(the_model.test_point))
import theano.tensor as tt
dgamma, dzeta, dxi = tt.grad(logl, [pm_gamma, pm_zeta, pm_xi])
inputdict = {pm_gamma: mle['x'][0], pm_zeta: mle['x'][1], pm_xi: mle['x'][2]}
print(logl.eval(inputdict), dgamma.eval(inputdict),
      dzeta.eval(inputdict), dxi.eval(inputdict))
with the_model:
    step = pm.NUTS()
q0 = step.potential.random()
p0 = step._logp_dlogp_func.dict_to_array(the_model.test_point)
print(p0, q0)
step.integrator.compute_state(q0, p0)<jupyter_output><empty_output><jupyter_text>Check where gradient breaks, and reparameterized<jupyter_code>loss = pm.math.log(qxus ** changes.astype(float)).sum()
dgamma, dzeta, dxi = tt.grad(loss, [pm_gamma, pm_zeta, pm_xi])
print(loss.eval(inputdict), dgamma.eval(inputdict),
      dzeta.eval(inputdict), dxi.eval(inputdict))
loss = tt.switch(changes, pm.math.log(qxus), 0.).sum()

dgamma, dzeta, dxi = tt.grad(loss, [pm_gamma, pm_zeta, pm_xi])
print(loss.eval(inputdict), dgamma.eval(inputdict),
      dzeta.eval(inputdict), dxi.eval(inputdict))
loss = pm.math.log(qxus[changes]).sum()

dgamma, dzeta, dxi = tt.grad(loss, [pm_gamma, pm_zeta, pm_xi])
print(loss.eval(inputdict), dgamma.eval(inputdict),
      dzeta.eval(inputdict), dxi.eval(inputdict))
x = tt.zeros(qxus.shape)
x = tt.inc_subtensor(x[changes], pm.math.log(qxus[changes]))
x = tt.inc_subtensor(x[~changes], pm.math.log(1.))

loss = x.sum()

dgamma, dzeta, dxi = tt.grad(loss, [pm_gamma, pm_zeta, pm_xi])
print(loss.eval(inputdict), dgamma.eval(inputdict),
      dzeta.eval(inputdict), dxi.eval(inputdict))<jupyter_output>140.2719692431427 451.11125744749626 217.28733567145977 308.8917065049014<jupyter_text>Implement a version that does not break the gradient<jupyter_code>def logp(gamma, zeta, xi):
    qxus = ((is_s * n_infected[:, None] * gamma) +
            (is_l * zeta) + (is_i * xi))
    x = tt.zeros(qxus.shape)
    x = tt.inc_subtensor(x[changes], qxus[changes])
    x = tt.inc_subtensor(x[~changes], 1.)

    logls = pm.math.log(x) - (qxus * dts[:, None])
    return pm.math.sum(logls)
print(-negative_log_likelihood(mle['x']))
print(logp(mle['x'][0], mle['x'][1], mle['x'][2]).eval())
loss = logp(pm_gamma, pm_zeta, pm_xi)
dgamma, dzeta, dxi = tt.grad(loss, [pm_gamma, pm_zeta, pm_xi])
print(loss.eval(inputdict), dgamma.eval(inputdict),
      dzeta.eval(inputdict), dxi.eval(inputdict))
with pm.Model() as the_model:
    gamma = pm.HalfNormal("gamma", sd=1)
    zeta = pm.HalfNormal("zeta", sd=1)
    xi = pm.HalfNormal("xi", sd=1)

    potential = pm.DensityDist("potential",
                               logp,
                               observed={"gamma": gamma,
                                         "zeta": zeta,
                                         "xi": xi})
the_model.test_point
the_model.check_test_point()
with the_model:
    mle = pm.find_MAP()
mle
with the_model:
    trace = pm.sample()
    
pm.traceplot(trace);<jupyter_output>Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (2 chains in 2 jobs)
NUTS: [xi, zeta, gamma]
Sampling 2 chains: 100%|██████████| 2000/2000 [00:56<00:00, 35.36draws/s]
{'fraction_non_alphanumeric': 0.16575258988421693, 'fraction_numerical': 0.03351614868982328, 'mean_word_length': 4.303715670436187, 'pattern_counts': {'":': 3, '<': 12, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '530185', 'license': 'MIT', 'n_tokens_mistral': 1372, 'n_tokens_neox': 1293, 'n_words': 224, 'path': 'PyMC3QnA/discourse_2772.ipynb', 'repo_name': 'junpenglao/Planet_Sakaar_Data_Science', 'size': '3282'}