tensorflow / probability

Probabilistic reasoning and statistical analysis in TensorFlow
https://www.tensorflow.org/probability/
Apache License 2.0
4.26k stars 1.1k forks source link

Composing two sigmoid for bijection gives nan log_prob #1556

Closed mlkrock closed 2 years ago

mlkrock commented 2 years ago

Hello, why does this calculation for log_prob return nan? The answer is -7280.1641, which I can calculate using other normalizing flow software, but tensorflow has some other things implemented I would like to use.

It also happens if I replace one of the Sigmoid with Exp or Square (or a more complicated bijector). And if i only use one Sigmoid bijector there is no issue. I'm not sure why elementwise composing Sigmoid with another bijector is a problem.

import tensorflow as  tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
sampledist = tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[2])
sigchain = tfb.Chain([tfb.Sigmoid(), tfb.Sigmoid()])
flow = tfd.TransformedDistribution(distribution=sampledist,bijector=sigchain)
flow.log_prob(tf.zeros(2))
csuter commented 2 years ago

Zero is not in the codomain of sigmoid(sigmoid(x)). The first sigmoid takes [-inf, inf] to [0, 1] and the second sigmoid takes [0, 1] to [.5, 1 / (1 + 1/e)] ~= [.5, .73].

Going backwards, the second sigmoid, inverted, takes zero to -inf, and the first sigmoid, inverted, will give garbage (eg nan) for this -inf input.

What software gives you a finite value for this? Can you share some code that successfully computes what you are expecting?

mlkrock commented 2 years ago

Thanks for the quick reply, sorry i got mixed up and thought it was the log_prob when input value was zero. Will close issue.

mlkrock commented 2 years ago

@csuter i won't reopen this because it's just a question no bug, but you might know quickly.

If i go on to write

tfk = tf.keras
tfkl = tfk.layers
x_ = tfkl.Input(shape=(2,), dtype=tf.float32)
log_prob_ = flow.log_prob(x_)

I get this long error with the logprob, presumably because of this tf.where in https://github.com/tensorflow/probability/blob/88d217dfe8be49050362eb14ba3076c0dc0f1ba6/tensorflow_probability/python/bijectors/sigmoid.py#L151

TypeError: You are passing KerasTensor(type_spec=TensorSpec(shape=(None, 2), dtype=tf.float32, name=None), name='tf.where/SelectV2:0', description="created by layer 'tf.where'"), an intermediate Keras symbolic input/output, to a TF API that does not allow registering custom dispatchers, such as `tf.cond`, `tf.function`, gradient tapes, or `tf.map_fn`. Keras Functional model construction only supports TF API calls that *do* support dispatching, such as `tf.math.add` or `tf.reshape`. Other APIs cannot be called directly on symbolic Kerasinputs/outputs. You can work around this limitation by putting the operation in a custom Keras layer `call` and calling that layer on this symbolic input/output.

Can you explain the work around, is it that I would need to directly define a "SigmoidSigmoid" custom layer?

csuter commented 2 years ago

Keras has its own tensor type that isn't compatible with many tf ops, which is the gist of what that error message is saying. If you want to use a distribution log prob as a loss, I think you can give a custom loss fun to keras compile (or fit? Not sure don't know keras that well). You can't pass an input layer directly in this way. Sorry, not sure if this is much help. Feel free to add more context on what you're trying to do, and I'll try to be more helpful.

On Wed, May 4, 2022 at 18:03 mlkrock @.***> wrote:

@csuter https://github.com/csuter i won't reopen this because it's just a question no bug, but you might know quickly.

If i go on to write

tfk = tf.keras tfkl = tfk.layers x_ = tfkl.Input(shape=(2,), dtype=tf.float32) logprob = flow.logprob(x)

I get this long error with the logprob, presumably because of this tf.where in https://github.com/tensorflow/probability/blob/88d217dfe8be49050362eb14ba3076c0dc0f1ba6/tensorflow_probability/python/bijectors/sigmoid.py#L151

TypeError: You are passing KerasTensor(type_spec=TensorSpec(shape=(None, 2), dtype=tf.float32, name=None), name='tf.where/SelectV2:0', description="created by layer 'tf.where'"), an intermediate Keras symbolic input/output, to a TF API that does not allow registering custom dispatchers, such as tf.cond, tf.function, gradient tapes, or tf.map_fn. Keras Functional model construction only supports TF API calls that do support dispatching, such as tf.math.add or tf.reshape. Other APIs cannot be called directly on symbolic Kerasinputs/outputs. You can work around this limitation by putting the operation in a custom Keras layer call and calling that layer on this symbolic input/output.

Can you explain the work around, is it that I would need to directly define a "SigmoidSigmoid" custom layer?

— Reply to this email directly, view it on GitHub https://github.com/tensorflow/probability/issues/1556#issuecomment-1117976742, or unsubscribe https://github.com/notifications/unsubscribe-auth/AABG2GMNOXLL46ZBHAX7EBDVILX3PANCNFSM5VC6GHDQ . You are receiving this because you were mentioned.Message ID: @.***>

mlkrock commented 2 years ago

@csuter Thanks I will give it a shot, but it's pretty similar and I'm afraid i am out of luck due to the Keras problem you described. i'm trying to fit data that is on the unit square [0,1] x [0,1] with a normalizing flow, which is why i keep asking about tfb.Sigmoid(). My hope is that I can compose the Sigmoid with some fancier bijector like tfb.MaskedAutoregressiveFlow so that I can be sure the fit lies in [0,1] x [0,1].

But i get this Keras error above, which I think is from checking tf.where in the sigmoid function.

I'm sort of following https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/AutoregressiveNetwork and https://simon-dirmeier.com/etudes/normalizing_flows.html as a guide for trying to compile this model, so here's my full example code.

import tensorflow as  tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
sampledist = tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[2])
sigchain = tfb.Chain([tfb.Sigmoid(), tfb.MaskedAutoregressiveFlow(tfb.AutoregressiveNetwork(params=2, hidden_units=[10, 10]))])
flow = tfd.TransformedDistribution(distribution=sampledist,bijector=sigchain)

tfk = tf.keras
tfkl = tfk.layers
x_ = tfkl.Input(shape=(2,), dtype=tf.float32)
log_prob_ = flow.log_prob(x_) #error here before i can compile the model and fit to data, same message as above
model = tfk.Model(x_, log_prob_)

Thanks for any help.