Improve numerical stability of censored logps

[aseyboldt]

Not using the more stable at.log1mexp instead of at.log(1 - at.exp(x)) can lead to numerical issues in the gradient, because the stabilization optimizations are applied after the gradient. eg something like:

import aesara
import aesara.tensor as at
import numpy as np

x = at.scalar("x")
a = at.scalar("a")
b = at.scalar("b")

logcdf = pm.logcdf(pm.Weibull.dist(a, b), x)
logccdf = at.log(1 - at.exp(logcdf))
logccdf2 = at.log1mexp(logcdf)

logcdf.name = "logcdf"
logccdf.name = "logccdf"

func = aesara.function([x, a, b], [at.grad(logccdf, a), at.grad(logccdf2, a)])

func(50, 1, 1)
# [array(-inf), array(-195.60115027)]

[codecov]

Codecov Report

Merging #156 (9bf8c24) into main (44e441f) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #156   +/-   ##
=======================================
  Coverage   94.88%   94.88%           
=======================================
  Files          12       12           
  Lines        1780     1780           
  Branches      263      263           
=======================================
  Hits         1689     1689           
  Misses         51       51           
  Partials       40       40           

Impacted Files	Coverage Δ
aeppl/truncation.py	98.27% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 44e441f...9bf8c24. Read the comment docs.

[brandonwillard]

Why wasn't this covered by a stabilizing rewrite?

[ricardoV94]

The only stabilizing rewrite we have is to introduce the log1mexp Op, but we don't have anything for the grad of the general graph of log1mexp. The grad of log1mexp was tweaked sometime ago here: aesara-devs/aesara#725

I am not sure we can or want to apply those tweaks to all graphs of that kind, or if they only work when we know they came from a grad of a log1mexp

It would help if grad was not eager and we could run some rewrites before introducing their graphs.