`JointDistributionCoroutineAutoBatched` seed behaviour with batches

Aperta
  • Qualcosa non va?
#2,011 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
35/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Ferma
Stack tecnologico
jupyter-notebook, python, tensorflow
Ambito
machine-learning

Direzione di ricerca

Inizia eseguendo il Colab collegato oppure l’esempio model_autobatched e le chiamate a sample_distributions con il backend JAX e campioni di theta vettorizzati. Traccia il campionamento e la gestione dei seed di JointDistributionCoroutineAutoBatched; il lavoro è completato quando i campioni indipendenti di y sono ottenuti lungo il batch di theta e in sample_distributions non si verifica alcuna conversione non supportata da object a Tensor.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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

Lingua principale
Jupyter Notebook
Stelle
4.4k
Fork
1.1k
Nessuna PR unita negli ultimi 30g

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di tensorflow/probability

Tutte le issue di tensorflow/probability

Issue simili

Altre issue su Machine Learning

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.