Error in HamiltonianMonteCarlo when using gather in log_prob function

Đang mở
  • Có gì không ổn?
#1,837 2 bình luận 0 reaction 0 người được giao Xem trên GitHub

Chưa có ai nhận issue này.

Đánh giá

Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức phù hợp với người mới
30/100
Loại issue
Lỗi
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Đình trệ
Công nghệ
python, tensorflow
Lĩnh vực
machine-learning

Hướng nghiên cứu

Bắt đầu với sample_hmc được cung cấp và việc tái hiện logprob, tập trung vào tfp.mcmc.sample_chain, HamiltonianMonteCarlo và ranh giới @tf.function. So sánh các trường hợp tf.gather và chỉ reduction, sau đó xác minh rằng lỗi cấu trúc được báo cáo đã được giải quyết, trong khi việc thực thi đồ thị vẫn được bật và sampling vẫn trả về các mẫu cùng các giá trị chấp nhận.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Mô tả

I am trying to run a Hamiltonian MCMC on a target distribution whose implementation involves a call to tf.gather. The following code:

import tensorflow as tf
import tensorflow_probability as tfp


def sample_hmc(
    target_log_prob_fn,
    current_state,
    num_results=1000,
    num_burnin_steps=500,
    adaptation_frac=0.8,
    num_leapfrog_steps=3,
    step_size=1.,
):    
    hmc = tfp.mcmc.SimpleStepSizeAdaptation(
        tfp.mcmc.HamiltonianMonteCarlo(
            target_log_prob_fn=target_log_prob_fn,
            num_leapfrog_steps=num_leapfrog_steps,
            step_size=step_size,
        ),
        num_adaptation_steps=int(num_burnin_steps * adaptation_frac))

    @tf.function
    def run_chain():
        samples, is_accepted = tfp.mcmc.sample_chain(
            num_results=num_results,
            num_burnin_steps=num_burnin_steps,
            current_state=current_state,
            kernel=hmc,
            trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
        )
        return samples, is_accepted
        
    return run_chain()


def logprob(alpha):
    indices = tf.constant([2, 0, 1], dtype=tf.int32)
    return -tf.math.reduce_sum(tf.gather(alpha**2, indices))
    # return -tf.math.reduce_sum(alpha**2)


alpha = tf.constant([1.0, 1.0, 1.0])
sample_hmc(
    logprob,
    current_state=alpha,
    num_results=10,
    num_burnin_steps=5,
)

raises ValueError: The two structures don't have the same nested structure. followed by a very long and (to me) cryptic message. Replacing the return statement in the logprob function with the commented line gets rid of the error. The error seems to appear whenever the result of logprob contains a tf.gather subexpression.

The error also disappears when I remove the @tf.function decorator from the definition of run_chain. However, this comes at a huge performance cost.

How can I efficiently sample from a distribution whose log-probability involves a tf.gather expression?

Ngôn ngữ chính
Jupyter Notebook
Star
4.4k
Fork
1.1k
Không có pull request nào được merge trong 30 ngày

Chuẩn bị môi trường

Bắt đầu từ đâu

  1. Đọc hết issue, rồi đọc hướng dẫn đóng góp của dự án.
  2. Bình luận trên issue rằng bạn sẽ nhận — tránh hai người làm cùng một việc.
  3. Fork repository và làm thay đổi trên một nhánh.
  4. Mở pull request có tham chiếu số hiệu của issue.

Issue khác của tensorflow/probability

Tất cả issue của tensorflow/probability

Issue tương tự

Thêm issue về Machine Learning

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.