每个类的tensorflow概率

import tensorflow_probability as tfp

tfd = tfp.distributions

class CustomDistribution(tfd.Distribution):
    def __init__(self, param1, param2, validate_args=False, allow_nan_stats=True, name='CustomDistribution'):
        parameters = dict(locals())
        super(CustomDistribution, self).__init__(dtype=tf.float32,
                                                 reparameterization_type=tfd.FULLY_REPARAMETERIZED,
                                                 validate_args=validate_args,
                                                 allow_nan_stats=allow_nan_stats,
                                                 name=name)
        self.param1 = param1
        self.param2 = param2

    def _log_prob(self, x):
        # 实现自定义的对数概率计算逻辑
        pass

surrogate_posterior = tfp.experimental.vi.build_factored_surrogate_posterior(
    event_shape=event_shape,
    bijector=[tfb.Softplus(), tfb.Softplus()]
)
losses = tfp.experimental.vi.fit_surrogate_posterior(
    target_log_prob_fn=target_log_prob_fn,
    surrogate_posterior=surrogate_posterior,
    optimizer=tf.optimizers.Adam(learning_rate=0.01),
    num_steps=num_steps
)

mcmc = tfp.mcmc.SimpleStepSizeAdaptation(
    inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target_log_prob_fn,
        num_leapfrog_steps=3,
        step_size=1.),
    num_adaptation_steps=int(num_burnin_steps * 0.8))

samples, kernel_results = tfp.mcmc.sample_chain(
    num_results=num_results,
    current_state=current_state,
    kernel=mcmc,
    num_burnin_steps=num_burnin_steps,
    parallel_iterations=1)