I'm currently learning Bayesian approaches through the book "Rethinking" and would like to utilize TensorFlow Probability. Specifically, I'm working on this chapter : https://colab.research.google.com/github/ksachdeva/rethinking-tensorflow-probability/blob/master/notebooks/04_geocentric_models.ipynb In this example, the model is defined using a function. However, I prefer to use the JointDistributionNamed approach as it closely aligns with how models are written in "Rethinking." I'm facing issues when running HMC, particularly with the target_log_prob_fn function. I'm unsure about the type of object returned within the HamiltonianMonteCarlo function.
Here is my code and the corresponding error message:
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd
tfb = tfp.bijectors
import pandas as pd
d = pd.read_csv('rethinking-master/data/Howell1.csv',sep = ';')
height= d.height
model = tfd.JointDistributionNamedAutoBatched(dict(
s = tfd.Sample(tfd.Exponential(1), sample_shape=1),
alpha = tfd.Sample(tfd.Normal(0,1), sample_shape=1),
beta = tfd.Sample(tfd.Normal(0,1), sample_shape=1),
y = lambda s,alpha,beta: tfd.Independent(tfd.Normal( alpha + beta * d.weight.values,s),
reinterpreted_batch_ndims=1),
))
def _trace_fn_transitioned(_, pkr):
return pkr.inner_results.inner_results.log_accept_ratio
num_chains = 4
num_leapfrog_steps = 4
step_size = 0.8
burnin = 500
params = ['s', 'alpha', 'beta']
init_state = list(model.sample(num_chains))[:-1]
bijectors = [tfb.Identity() for i in init_state]
observed_data=(d.height.values)
def target_log_prob_fn (**x):
print(x[0])
return model.log_prob(model.sample(y = observed_data, **x))
hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn, num_leapfrog_steps=num_leapfrog_steps, step_size=step_size
)
inner_kernel = tfp.mcmc.TransformedTransitionKernel(
inner_kernel=hmc_kernel, bijector=bijectors
)
kernel = tfp.mcmc.SimpleStepSizeAdaptation(
inner_kernel=inner_kernel,
target_accept_prob=0.8,
num_adaptation_steps=int(0.8 * burnin),
log_accept_prob_getter_fn=lambda pkr: pkr.inner_results.log_accept_ratio,
)
tfp.mcmc.sample_chain(
num_results=544,
num_burnin_steps=burnin,
current_state=init_state,
kernel=kernel,
trace_fn=_trace_fn_transitioned,
)
error: target_log_prob_fn() takes 0 positional arguments but 3 were given
Thanks,
To use the unpacking approach with
**x
, you would need to adjust how the HMC kernel passes the arguments. Unfortunately, this isn't straightforward because the TFP MCMC framework is designed to work with positional arguments, and I have no idea how to help you there.The standard way to define target_log_prob_fn would be
Moreover, there's an issue with how you're sampling initial states and creating bijectors. Since you have three parameters (
s, alpha, beta
), your initial state should be a list of tensors, each corresponding to one of these parameters. The bijectors should match these parameters as well. You can do it like thisRunning it on my machine has the desired outcome, so I believe that it should be the same for you. Although it may not be the reply you were looking for, I hope it stil helps