无法使用TensorFlow_probability api采样某些概率（TypeError：__init __（）缺少1个必需的位置参数：'value_index'）

Question

我想使用TensorFlow_probability api进行一些采样。以二维高斯分布为例。我找不到很好的教程。我从TensorFlow网站修改的代码不起作用。

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

tf.enable_eager_execution()
def unnormalized_log_prob(x): # Two dimensional unnormalized Gaussian
    #dim x= batch size,vec size?
    kernel = tf.Tensor([[1.,0.],[0.,10.]], dtype = tf.float32)
    mu = tf.Tensor([1.,5.], dtype = tf.float32)
    return -0.5*tf.matmul((x-mu), tf.matmul(tf.inverse(kernel), tf.transpose(x-mu)))

# Initialize the HMC transition kernel.
num_results = int(10e3)
num_burnin_steps = int(1e3)
adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
    tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=unnormalized_log_prob,
        num_leapfrog_steps=3,
        step_size=1.),
    num_adaptation_steps=int(num_burnin_steps * 0.8))

# Run the chain (with burn-in).
@tf.function
def run_chain():
  # Run the chain (with burn-in).
    samples, is_accepted = tfp.mcmc.sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=np.array([[1.,1.],[3.,3.],[5.,5.]], dtype=np.float32),
        kernel=adaptive_hmc,
        trace_fn=lambda _, pkr: pkr.inner_results.is_accepted)

    sample_mean = tf.reduce_mean(samples)
    sample_stddev = tf.math.reduce_std(samples)
    is_accepted = tf.reduce_mean(tf.cast(is_accepted, dtype=tf.float32))
    return samples, sample_mean, sample_stddev, is_accepted

samples, sample_mean, sample_stddev, is_accepted = run_chain()

print('mean:{:.4f}  stddev:{:.4f}  acceptance:{:.4f}'.format(
    sample_mean.numpy(), sample_stddev.numpy(), is_accepted.numpy()))

我希望如果我想分别从初始状态[1.，1。]和[3.，3。]，[5.，5。]采样三个样本链，它将返回维度为[ 3，10000,2]