fix the case that value < 0 for exp and geo dist

[blacksde]

resolve #1141

[blacksde]

@brianwa84 could you help review this issue? Thx~

[axch]

LGTM.

[csuter]

My default is not to make this change. 1) We often permit computation of analytical pdf/cdf formulas outside the support (eg, bernoulli or poisson at a non-integer). 2) We also generally default to not adding compute/RAM overhead by default, which "punishes" users who have already taken care not to feed inputs outside the support (eg, by constraining inputs, casting, etc). "Punish" here refers to the fact that with this change, all users will incur a computational/RAM cost (tf.where will cause both sides of the expression to be evaluated). We could save RAM in this change by not using zeros_like, but instead a scalar tf.constant(0., dtype=value.dtype), which uses only a fixed amount of RAM, but we're still creating the where op which adds to, eg, graph/xla compilation time; these little changes can really add up if we're not vigilant. Note that when validate_args is True (off by default, again to save compute overhead), an assertion will be triggered on values outside the support.

[axch]

While I respect the point about compute and compilation cost, I think there is at this point strong evidence that our users expect Exponential.cdf to compute the CDF of the Exponential distribution, not an analytic continuation of the CDF of the Exponential distribution from the region (0, +inf) to R. We have had multiple bug reports and at least one PR about this kind of thing.

Perhaps we should add a new flag to control these checks? It could be named analytic_extensions=False or avoid_tf_where=False, depending on whether we wish to emphasize the performance point or the semantic point. I propose setting the flag to False by default (i.e., including the tf.where by default) because that is safest for users starting to learn our system; and users who need the extra performance and wish to accept the extra responsibility can change the flag.

[blacksde]

@csuter Ths for your clarification of the initial idea of the design, but we were still hoping that we could have this formula mathematically right. As explained in #1142 , the current formula will lead to some unreasonable numbers such as prob < 0 or > 1.

@axch Thx for helping review this pr. Let me know if there is any further work from me(it looks that some ci failed)

[axch]

@blacksde After some internal discussion, we decided to add a generically-named flag to control whether to have these kinds of tf.where operations.

Rationale:

• There are situations, such as your use-case, where mathematical correctness over all of R is valuable.
• On the other hand, it's generally not too hard to set up unconstraining bijectors in such a way that the distribution is never queried for a log_prob etc outside of its support. For those cases, it was deemed valuable to permit these bounds checks to be elided from the TensorFlow graph entirely (which, largely, corresponds to the status quo).

Action requested:

• Please add a constructor flag to Geometric and Exponential named force_probs_to_zero_outside_support, default value False. Please make sure that when the flag is False, the TF graph consists of the computation as it was before your change, and only include the tf.where you are adding here when the flag is True.
• See Poisson for an example.

[blacksde]

@blacksde After some internal discussion, we decided to add a generically-named flag to control whether to have these kinds of tf.where operations.

Rationale:

• There are situations, such as your use-case, where mathematical correctness over all of R is valuable.
• On the other hand, it's generally not too hard to set up unconstraining bijectors in such a way that the distribution is never queried for a log_prob etc outside of its support. For those cases, it was deemed valuable to permit these bounds checks to be elided from the TensorFlow graph entirely (which, largely, corresponds to the status quo).

Action requested:

• Please add a constructor flag to Geometric and Exponential named force_probs_to_zero_outside_support, default value False. Please make sure that when the flag is False, the TF graph consists of the computation as it was before your change, and only include the tf.where you are adding here when the flag is True.
• See Poisson for an example.

The plan looks great~Let me work on the update.

[blacksde]

@axch @csuter I have updated the code according to the comments. Could you help review the code?Thx~