Can Jax generate numerically stable functions?

[NeilGirdhar]

Is it possible for Jax to take a function, calculate a set of derivatives, produce the Taylor expansion, and then use that to produce a numerically stable version of the function?

[mattjj]

Cool question! I don't know :) Got any references for the technique/algorithm you have in mind?

[NeilGirdhar]

I don't. The reason I thought of this is that I noticed that scipy.stats.beta produces a more accurate entropy than my efax.BetaNP.entropy—even though my function is analytical. Then I saw that they have special cased the entropy function for extreme values. I feel like it should be possible for Jax to somehow generate those accurate functions for extreme values?

[ArvinSKushwaha]

This project: https://herbie.uwplse.org/ may be an interesting reference?

[NeilGirdhar]

@ArvinSKushwaha Yes, very cool! So this is possible in some cases at least. Thanks for sharing.