Error in HamiltonianMonteCarlo when using gather in log_prob function

オープン
  • 何か問題がありますか？
#1,837 コメント 2 件 リアクション 0 件 担当者 0 名 GitHub で見る

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
30/100
issue の種類
バグ
明瞭さ
おおむね明確
活発さ
停滞
技術スタック
python, tensorflow
領域
machine-learning

調査の方向性

提供された sample_hmc と logprob の再現から始め、tfp.mcmc.sample_chain、HamiltonianMonteCarlo、@tf.function の境界に焦点を当てます。tf.gather のケースと reduction のみのケースを比較し、グラフ実行を有効にしたまま、報告された構造エラーが解消され、サンプリングが引き続きサンプルと受理値を返すことを確認します。

索引モデルが issue の本文から書いたものです。

説明

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?

主要言語
Jupyter Notebook
スター
4.4k
フォーク
1.1k
30日以内にマージされた PR はありません

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

tensorflow/probability のほかの issue

tensorflow/probability の issue をすべて見る

似ている issue

Machine Learning の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。