def ais_step(pot_pre, pot_new, sample_pre, log_weight_pre, mcmc_kernel):
    """
    Compute log Z(p*) = \sum log_avg_exp(log w_T(x_i)) + log Z(p_0).
    It includes two steps:
    1) update w(x_i) by the ratio of two adjacent annealing distribution
    log w_{t}(x_i^t) = log w_{t-1}(x_i^{t-1}) + log p^*_{t}(x_i^{t-1}) - log p^*_{t-1}(x_i^{t-1})
    2) update the sample by HMC transition
    x_i^t ~ T_{HMC}(x, x_{i}^{t-1}).

    :param pot_pre: the potential function of p_{t-1}
    :param pot_new: the potential function of p_{t}(x_i^{t-1})
    :param sample_pre: the batch of samples x^{t-1}
    :param log_weight_pre: the batch of previous log weights log w_{t-1}(x^{t-1})
    :param mcmc_kernel: MCMC kernel for updating sample batch x_{t-1}
    :return: sample_new: updated sample batch x_{t}
             log_weight_new: updated log weight batch log w_{t}(x^t)
             acp_flag: boolean array for acceptance of new samples
             w_update: the log weight update log p^*_{t}(x_i^{t-1}) - log p^*_{t-1}(x_i^{t-1})
    """
    w_update = -pot_new(sample_pre) + pot_pre(sample_pre)  # log f_t(x_{t-1}) - log f_{t-1}(x_{t-1})
    log_weight_new = log_weight_pre + w_update
    sample_new, _, _, acp_flag = mcmc_kernel(pot_new, sample_pre)
    return sample_new, log_weight_new, acp_flag, w_update