`JointDistributionCoroutineAutoBatched.sample_distributions` errors using jax substrate

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

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

評価

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

調査の方向性

まず、提供されている JointDistributionCoroutineAutoBatched.sample_distributions の JAX 再現コードを実行し、動作する sample のケースおよび自動バッチ処理を使わないケースと比較します。sample_distributions のエントリポイントと、サポートされていないオブジェクトから Tensor への変換を追跡します。自動バッチ処理された JAX 呼び出しが ValueError なしで分布とサンプルを返せば完了です。

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

説明

Hi, I found that sample_distributions can error when using the jax substrate. sample seems to be ok, using tf appears to be ok, and using the non auto batched joint distribution is also ok.

Here is a small example:

from functools import partial

import jax
from tensorflow_probability.substrates import jax as tfp

tfd = tfp.distributions

tfp.__version__
# 0.25.0

@partial(tfd.JointDistributionCoroutine, batch_ndims=0)
def joint_dist():
    x = yield tfd.Gamma(2.0, 10.0, name="x")
    y = yield tfd.Gamma(x, 10.0, name="y")


seed = jax.random.key(123)

# ok
dists, samples = joint_dist.sample_distributions(x=[1.0, 2.0], seed=seed)

# samples:
# StructTuple(
#   x=Array([1., 2.], dtype=float32),
#   y=Array([0.10665689, 0.21802416], dtype=float32)
# )


@tfd.JointDistributionCoroutineAutoBatched
def joint_dist():
    x = yield tfd.Gamma(2.0, 10.0, name="x")
    y = yield tfd.Gamma(x, 10.0, name="y")


# ok
samples = joint_dist.sample(x=[1.0, 2.0], seed=seed)

# samples:
# StructTuple(
#   x=Array([1., 2.], dtype=float32),
#   y=Array([0.05508393, 0.14792603], dtype=float32)
# )

# ValueError: Attempt to convert a value (<object object at 0x717aa57398a0>) with an unsupported type (<class 'object'>) to a Tensor.
dists, samples = joint_dist.sample_distributions(x=[1.0, 2.0], seed=seed)

Thanks!

主要言語
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 を短くまとめたダイジェスト。