log_prob of blocked programs

[jhn-nt]

Hello,

First, thanks for developing such an amazing package. I am newbie to oryx and was playing around with its functionalities, perhaps naively I had been attempting to evalute log_probs of blocker porblems, as below:

from jax.random import split
from oryx.core import ppl
import tensorflow_probability.substrates.jax.distributions as tfd

def latent_normal(key):
    z_key,x_key= split(key)
    z=ppl.random_variable(tfd.Normal(0,1),name="z")(z_key)
    return ppl.random_variable(tfd.Normal(z,1e-1),name="x")(x_key)


blocked=ppl.block(latent_normal,names=["z"])
ppl.joint_log_prob(blocked)({"x":10})

However, it returns: { "name": "ValueError", "message": "Cannot compute log_prob of function.", "stack": "--------------------------------------------------------------------------- ValueError Traceback (most recent call last) Cell In[2], line 12 8 return ppl.random_variable(tfd.Normal(z,1e-1),name="x")(x_key) 11 blocked=ppl.block(latent_normal,names=["z"]) ---> 12 ppl.joint_log_prob(blocked)({"x":10})

File /opt/conda/lib/python3.11/site-packages/oryx/core/interpreters/log_prob.py:71, in log_prob..wrapped(sample, *args, **kwargs) 67 flat_incells = [ 68 InverseAndILDJ.unknown(trace_util.get_shaped_aval(dummy_seed)) 69 ] + [InverseAndILDJ.new(val) for val in flat_inargs] 70 flat_outcells = [InverseAndILDJ.new(a) for a in flat_outargs] ---> 71 return log_prob_jaxpr(jaxpr.jaxpr, constcells, flat_incells, flat_outcells)

File /opt/conda/lib/python3.11/site-packages/oryx/core/interpreters/log_prob.py:128, in log_prob_jaxpr(jaxpr, constcells, flat_incells, flat_outcells) 118 _, final_log_prob = propagate.propagate( 119 InverseAndILDJ, 120 log_prob_rules, (...) 125 reducer=reducer, 126 initial_state=0.) 127 if final_log_prob is failed_log_prob: --> 128 raise ValueError('Cannot compute log_prob of function.') 129 return final_log_prob

ValueError: Cannot compute log_prob of function." }

Am I missing something?

Thanks again

Very Best Giovanni

[PaulScemama]

Hi @jhn-nt. I'm not a contributor or maintainer but I think you might be missing the following. If you look at this description of block it says:

"The block transformation takes in a program and a sequence of names and returns a program that behaves identically except that in downstream transformations (likejoint_sample), the provided names are ignored."

When you run blocked=ppl.block(latent_normal,names=["z"]) you are telling downstream transformations (like joint_log_prob) to ignore the variable with name 'z'. Because the variable with name 'x' depends on the variable with name 'z', which you blocked, joint_log_prob is unable to "see" what the value of 'z' is to get the distribution over 'x' to then compute the log_prob of 10 under that distribution. That's why you're seeing the error.

In general, I would guess you are unable to compute samples or log probabilities of a program with respect to a variable that depends on variables you previously blocked. @sharadmv let me know if I missed anything; I am not sure if there is a principled way to catch this sort of error and display a more helpful error message -- I may look into it. And @jhn-nt let me know if you have any further questions! Hope that helps.

[jhn-nt]

Thank very much! this really helped

Giovanni