`JointDistributionCoroutineAutoBatched` seed behaviour with batches

Abierto
  • ¿Algo no está bien?
#2,011 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
35/100
Tipo de issue
Error
Claridad
Bastante claro
Estado de actividad
Estancado
Stack tecnológico
jupyter-notebook, python, tensorflow
Área
machine-learning

Línea de trabajo

Empieza ejecutando el Colab enlazado o el ejemplo model_autobatched y las llamadas a sample_distributions con el backend JAX y muestras de theta vectorizadas. Rastrea el muestreo y la gestión de seeds de JointDistributionCoroutineAutoBatched; se considera terminado cuando haya extracciones independientes de y a través del batch de theta y no se produzca ninguna conversión no compatible de object a Tensor en sample_distributions.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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

Lenguaje dominante
Jupyter Notebook
Estrellas
4.4k
Forks
1.1k
Sin PR fusionados en 30 d

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de tensorflow/probability

Todos los issues de tensorflow/probability

Issues similares

Más issues de Machine Learning

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.