`JointDistributionCoroutineAutoBatched` seed behaviour with batches

Ouverte
  • Quelque chose ne va pas ?
#2,011 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub

Personne n'a encore pris cette issue.

Évaluation

Difficulté
4/5
Temps estimé
3-5 jours
Accessibilité débutants
35/100
Type d'issue
Bug
Clarté
Plutôt claire
Activité
À l'abandon
Stack technique
jupyter-notebook, python, tensorflow
Domaine
machine-learning

Piste de recherche

Commencez par exécuter le Colab lié ou l’exemple model_autobatched ainsi que les appels à sample_distributions avec le backend JAX et des échantillons de theta vectorisés. Suivez l’échantillonnage et la gestion des seeds de JointDistributionCoroutineAutoBatched ; le travail est terminé lorsque les tirages de y sont indépendants sur le batch de theta et qu’aucune conversion non prise en charge de object vers Tensor ne se produit dans sample_distributions.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Description

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

Langage dominant
Jupyter Notebook
Étoiles
4.4k
Forks
1.1k
Aucune PR mergée en 30 j

Préparer son environnement

Par où commencer

  1. Lisez l'issue en entier, puis le guide de contribution du projet.
  2. Signalez en commentaire que vous la prenez — cela évite que deux personnes fassent le même travail.
  3. Forkez le dépôt et travaillez sur une branche.
  4. Ouvrez une pull request qui référence le numéro de l'issue.

Autres issues de tensorflow/probability

Toutes les issues de tensorflow/probability

Issues similaires

Plus d'issues Machine Learning

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.