Is there an efficient way to generate MCMC samples repeatedly?

Hi, can someone advise on proper way (vectorize?) to generate MCMC samples repeatedly?

I have many (e.g., 10,000) sets of data and need to run MCMC samples for each data set to get summary statistics from posterior distributions. I can use a for loop for the process but it may take a long time if each MCMC process takes some time. I wonder if there is a proper tf/tfp way to do that.

Here is the loop way modifying the example from https://www.tensorflow.org/probability/examples/A_Tour_of_TensorFlow_Probability. It ran ok but takes time in multiple of the data set number (num_datasets). Wonder if there is a way to absorb datasets in some batch dimension to avoid the for loop. Thanks!

import tensorflow_probability as tfp
import tensorflow as tf
import numpy as np

tfd = tfp.distributions

generate num_datasets data sets

num_datasets=100 ### number of data sets
num_features = 2
num_examples = 50
noise_scale = .5
true_w = np.array([-1., 2., 3.])

distribution with mean as true_w with some variance

dist_w = tfd.MultivariateNormalDiag(
loc=true_w,
scale_diag=[0.3, 0.4,0.5])

def f(x, w):

Pad x with 1's so we can add bias via matmul

x = tf.pad(x, [[1, 0], [0, 0]], constant_values=1)
linop = tf.linalg.LinearOperatorFullMatrix(w[..., np.newaxis])
result = linop.matmul(x, adjoint=True)
return result[..., 0, :]

sample_w=dist_w.sample(num_datasets)
sample_xs=np.random.uniform(-1., 1., [num_datasets,num_features, num_examples])

sample_ys=np.empty([num_datasets, num_examples])
for n in range(num_datasets):
sample_ys[n,:]=f(sample_xs[n,:], sample_w[n,:].numpy()) + np.random.normal(0., noise_scale, size=num_examples)

Define the joint_log_prob function, and our unnormalized posterior.

def joint_log_prob(w, x, y):

Our model in maths is

w ~ MVN([0, 0, 0], diag([1, 1, 1]))

y_i ~ Normal(w @ x_i, noise_scale), i=1..N

rv_w = tfd.MultivariateNormalDiag(
loc=np.zeros(num_features + 1),
scale_diag=np.ones(num_features + 1))

rv_y = tfd.Normal(f(x, w), noise_scale)
return (rv_w.log_prob(w) +
tf.reduce_sum(rv_y.log_prob(y), axis=-1))

r_list=[]

import time
start_time = time.time()

for n in range(num_datasets):

get each data set

xs = sample_xs[n,:]
ys = sample_ys[n,:]

Create our unnormalized target density by currying x and y from the joint.

def unnormalized_posterior(w):
return joint_log_prob(w, xs, ys)

# Create an HMC TransitionKernel

hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=unnormalized_posterior,
step_size=np.float64(.1),
num_leapfrog_steps=2)

We wrap sample_chain in tf.function, telling TF to precompile a reusable

computation graph, which will dramatically improve performance.

@tf.function
def run_chain(initial_state, num_results=1000, num_burnin_steps=500):
return tfp.mcmc.sample_chain(
num_results=num_results,
num_burnin_steps=num_burnin_steps,
current_state=initial_state,
kernel=hmc_kernel,
trace_fn=lambda current_state, kernel_results: kernel_results)

initial_state = np.zeros(num_features + 1)
samples, kernel_results = run_chain(initial_state)

print("Acceptance rate:", kernel_results.is_accepted.numpy().mean())

r_samples=tfp.stats.percentile(samples, q=3., axis=0)
r_list.append(r_samples[tf.newaxis,...])
print(n, "--- %s seconds ---" % (time.time() - start_time))

print('Result from all data',tf.concat(r_list,axis=0).shape)

Start with the provided joint_log_prob function, the per-dataset loop, and the tfp.mcmc.sample_chain call. Investigate whether the existing MCMC API supports batching these datasets without the loop; done means a validated approach that produces the requested posterior summaries for all datasets.