Using JointDistribution class. Is this a bug?

[Ashutosh Tewari]

Hello TFP community,

I am trying to code up a fairly simple joint distribution p(x,y), defined as follows

x~Bernoulli(p_success=0.7)

y~Normal(0,1) if x else Laplace(0,1)

I believe this can be accomplished using the JointDistribution APIs such JointDistributionNamed. Following is the code I am using.

@tf.function
def condDist(e):
    return tfd.Normal(loc=0.0,scale=1.0) if e else tfd.Laplace(loc=0.0,scale=1.0)
   
joint = tfd.JointDistributionNamed(
     dict(x = tfd.Bernoulli(probs=0.7),
          y = lambda x: condDist(x)
     ),
     batch_ndims=0,
     use_vectorized_map=True)

I am able to define this joint distribution, but getting an error when calling the sample() function (see the error message at the end). Apparently the tf.cond() function (that encodes the IF statement in the conditional distribution) doesn't like two different distributions (Normal vs. Laplace) are outputted based on whether x is 0 or 1. This should certainly be permissible, hence my hunch is that it is some sort of bug. Any insight will be greatly appreciated.

File "/tmp/ipykernel_37/3096765418.py", line 3, in condDist *

return tfd.Normal(loc=0.0,scale=1.0) if e else tfd.Laplace(loc=0.0,scale=1.0) TypeError: true_fn and false_fn arguments to tf.cond must have the same number, type, and overall structure of return values. true_fn output: tfp.distributions.Normal("Normal_1_1", batch_shape=[], event_shape=[], dtype=float32) false_fn output: tfp.distributions.Laplace("Laplace_1_1", batch_shape=[], event_shape=[], dtype=float32) Error details: The two structures don't have the same nested structure. First structure: type=Normal str=tfp.distributions.Normal("Normal_1_1", batch_shape=[], event_shape=[], dtype=float32) Second structure: type=Laplace str=tfp.distributions.Laplace("Laplace_1_1", batch_shape=[], event_shape=[], dtype=float32) More specifically: Incompatible CompositeTensor TypeSpecs: type=Normal_ACTTypeSpec str=Normal_ACTTypeSpec(3, {'loc': TensorSpec(shape=(), dtype=tf.float32, name=None), 'scale': TensorSpec(shape=(), dtype=tf.float32, name=None)}, {'validate_args': False, 'allow_nan_stats': True, 'name': 'Normal_1_1'}, ('parameters',), (), ('name',), {}) vs. type=Laplace_ACTTypeSpec str=Laplace_ACTTypeSpec(3, {'loc': TensorSpec(shape=(), dtype=tf.float32, name=None), 'scale': TensorSpec(shape=(), dtype=tf.float32, name=None)}, {'validate_args': False, 'allow_nan_stats': True, 'name': 'Laplace_1_1'}, ('parameters',), (), ('name',), {})

[Colin Carroll]

Hi Ashutosh - 

Conditionals can be tricky! To create the model you are asking about, I would use

```

@tf.function
def condDist(e):
  e = tf.cast(e, tf.float32)
  return tfd.Mixture(cat=tfd.Categorical(probs=[1. - e, e]),
    components=[
      tfd.Normal(loc=0., scale=1.),
      tfd.Laplace(loc=0., scale=1.)])

   
joint = tfd.JointDistributionNamed(
     dict(x = tfd.Bernoulli(probs=0.7),
          y = lambda x: condDist(x)
     ),
     batch_ndims=0,
     use_vectorized_map=True)

```

Note that if you don't actually care about the value of `x`, you can directly use the `tfd.Mixture` above, with `e = 0.7`.

--
You received this message because you are subscribed to the Google Groups "TensorFlow Probability" group.
To unsubscribe from this group and stop receiving emails from it, send an email to tfprobabilit...@tensorflow.org.
To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/tfprobability/c49c3152-ca4c-4a8f-bb53-bfdbb773f392n%40tensorflow.org.

[Ashutosh Tewari]

Colin- Many Thanks for your prompt response. Initially I found employing tfd.Mixture() to specify the conditional in this manner to be a bit convoluted. However, the more I think about it the more I am convinced this may be a good way to specify distributions conditioned on discrete random variables. The tfd.Mixture() implements the more intuitive IF-ELSE conditions more compactly. Greatly appreciate your response.   BTW, when calling the joint.sample() I getting a warning shown below. Any clue what's it about? What's would be a good way to suppress it?

2023-06-20 18:52:14.909719: I tensorflow/core/common_runtime/executor.cc:1197] [/device:CPU:0] (DEBUG INFO) Executor start aborting (this does not indicate an error and you can ignore this message): INVALID_ARGUMENT: You must feed a value for placeholder tensor 'loop_body/iid_sample_fn_stateful_body/make_rank_polymorphic/fn_of_vectorized_args/Bernoulli/sample/uniform/stateless_random_uniform/StatelessRandomUniformV2/pfor/while/loop_body/iid_sample_fn_stateful_body/make_rank_polymorphic/fn_of_vectorized_args/Bernoulli/sample/uniform/stateless_random_uniform/shape' with dtype int32 and shape [1] [[{{node loop_body/iid_sample_fn_stateful_body/make_rank_polymorphic/fn_of_vectorized_args/Bernoulli/sample/uniform/stateless_random_uniform/StatelessRandomUniformV2/pfor/while/loop_body/iid_sample_fn_stateful_body/make_rank_polymorphic/fn_of_vectorized_args/Bernoulli/sample/uniform/stateless_random_uniform/shape}}]]