`JointDistributionCoroutineAutoBatched` seed behaviour with batches

Đang mở
#2,011 0 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
35/100
Loại issue
Lỗi
Độ rõ ràng
Khá rõ ràng
Mức độ hoạt động
Đình trệ
Công nghệ
jupyter-notebook, python, tensorflow
Lĩnh vực
machine-learning

Hướng nghiên cứu

Bắt đầu bằng cách chạy Colab được liên kết hoặc sample model_autobatched và các lệnh gọi sample_distributions với backend JAX và các mẫu theta được vector hóa. Truy vết quá trình lấy mẫu và xử lý seed của JointDistributionCoroutineAutoBatched; hoàn thành khi các lần lấy y là độc lập trên batch theta và không có chuyển đổi object-to-Tensor không được hỗ trợ trong sample_distributions.

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

Mô tả

Hi all,

I'm trying to compute a posterior predictive distribution over samples from a posterior distribution (Colab here). TFP 0.25 with JAX backend.

My (mre and therefore contrived) model specification is

@tfd.JointDistributionCoroutineAutoBatched
def model_autobatched():
    theta = yield tfd.Normal(loc=0., scale=1., name="theta")
    yield tfd.Normal(loc=theta, scale=0.1, name="y")

i.e. a Normally-distributed observation model with Normally-distributed mean. To compute the posterior predictive distribution, I wish to sample the y component conditional on a vector of theta samples.

theta_samples = np.arange(5.)
model_autobatched.sample(theta=theta_samples, seed=jax.random.key(0))

giving

StructTuple(
  theta=Array([0., 1., 2., 3., 4.], dtype=float32),
  y=Array([0.06215769, 1.0621576 , 2.0621576 , 3.0621576 , 4.0621576 ],      dtype=float32)
)

Oh dear, we notice that y - theta = constant. This seems to suggest that a single PRNG key is being used for each draw of y given the sample from theta.

Moreover, this approach fails entirely if sample_distributions is called.

model_autobatched.sample_distributions(theta=theta_samples, seed=jax.random.key(0))
ValueError: Attempt to convert a value (<object object at 0x7a53561590d0>) with an unsupported type (<class 'object'>) to a Tensor.

As a workaround, we could use the older JointDistributionCoroutine with Root annotation which works as desired (see Colab)

[edit] actually, JDCoroutine/Root only works because the whole theta vector is passed to y's constructor, not vectorisation over the whole model.

Do we have a bug or a feature, I wonder?

Regards,

Chris

Ngôn ngữ chính
Jupyter Notebook
Star
4.4k
Fork
1.1k
Chỉ số merge pull request
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.