sampling from MultivariateNormalDiag with placeholders - tensorflow2.0

Question

I'm running tensorflow 2.1 and tensorflow_probability 0.9. I'd like to parameterize a multivariatenormal distribution with placeholders and sample from it. Here's what I've tried
import tensorflow as tf
import tensorflow_probability as tfp
#tf.function()
def sample_vae(dist):
return dist.sample()
vae_mu = tf.keras.layers.Input(shape=(5), dtype=tf.float16)
vae_logvar = tf.keras.layers.Input(shape=(5), dtype=tf.float16)
dist = tfp.distributions.MultivariateNormalDiag(loc=vae_mu, scale_diag=tf.exp(vae_logvar))
z = sample_vae(dist)
The above gives me the following error
TypeError Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
60 op_name, inputs, attrs,
---> 61 num_outputs)
62 except core._NotOkStatusException as e:
TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
#tf.function
def has_init_scope():
my_constant = tf.constant(1.)
with tf.init_scope():
added = my_constant * 2
The graph tensor has name: BroadcastArgs_3:0
During handling of the above exception, another exception occurred:
_SymbolicException Traceback (most recent call last)
6 frames
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
73 raise core._SymbolicException(
74 "Inputs to eager execution function cannot be Keras symbolic "
---> 75 "tensors, but found {}".format(keras_symbolic_tensors))
76 raise e
77 # pylint: enable=protected-access
_SymbolicException: Inputs to eager execution function cannot be Keras symbolic tensors, but found [<tf.Tensor 'BroadcastArgs_3:0' shape=(1,) dtype=int32>, <tf.Tensor 'Exp_3:0' shape=(None, 5) dtype=float16>, <tf.Tensor 'input_7:0' shape=(None, 5) dtype=float16>]

Answer 1

I was able to get around this using DistributionLambda
vae_mu = tf.keras.layers.Input(shape=(1, 5), dtype=tf.float16)
vae_logvar = tf.keras.layers.Input(shape=(1, 5), dtype=tf.float16)
mu_logvar = tf.concat([vae_mu, vae_logvar], axis=1)
l = vae_mu.shape[1]
dist = tfp.layers.DistributionLambda(lambda theta: tfp.distributions.MultivariateNormalDiag(loc=theta[:, :l], scale_diag=theta[:, l:]))
z = dist(mu_logvar)