HI @Majid_Mohammadi – the error you see is what the commit I mentioned should fix. I think you just tried it a little too soon. The new tfp-nightly package goes out around 08:00 UTC. I think you will also need a couple of amendments to the code I wrote:
    w = yield tfd.Dirichlet(concentration=beta_agg[..., tf.newaxis] * w_star, name=f'w{i}')

will make sure the shapes align correctly. And you’ll also want to wrap the hmc call with tf.function:
results = tf.function(lambda: tfp.experimental.mcmc.windowed_adaptive_hmc(
    n_draws=100,
    joint_dist=pinned,
    num_leapfrog_steps=10,
    n_chains=100,
    num_adaptation_steps=300))()

This will ensure the sampler is fully compiled and should reduce run time.