Improve the partial derivatives of tfp_math.betainc by leandrolcampos · Pull Request #1580 · tensorflow/probability (original) (raw)

This PR tries to improve the partial derivatives of tfp_math.betainc(a, b, x) with respect to a and b.

Some of the improvement opportunities I identified in the current implementation of these partial derivatives are:

When dtype is float64:
- Moderate loss of accuracy when a and b are uniformly distributed in the interval (0, 10).
- When a and b are very small, gradients contain a very high proportion of NaN.
When dtype is float32:
- Heavy loss of accuracy in all domains evaluated.
- Gradients contain NaN (sometimes infinite as well) in all domains evaluated.
- When a and b are very small, gradients contain only NaN.

The alternative implementation I propose addresses all these improvement opportunities. Please see the following notebook:
https://gist.github.com/leandrolcampos/776048810f9c4ae18fc9598f54543057