Add incomplete gamma functions

i.e., scipy.special.gammainc and scipy.special.gammaincc.

[mattjj]

Awesome, thanks for your interest!

AIUI there are basically two ways to add functions like these:

1. implement them ourselves in terms of lax primitives, like how we implement, say, expit or logaddexp
2. find an implementation in the C++ XLA client math library and add it to JAX as a primitive in lax.py, like we do with lax.digamma here.

Interestingly, these two approaches are pretty similar: the XLA client libraries just implement functions in terms of XLA:HLO, and that’s exactly what we get when we write functions at the Python level in terms of the primitives in lax.py. That’s because the lax.py primitives are essentially 1:1 with XLA HLO (or in a few cases client library functions), so they translate to the same XLA computations. In general I think it’s easier and nicer to implement things in Python in terms of lax.py, but on the other hand for trickier numerical computations it can be good to reuse the effort in the XLA client library, and then we also get to share any updates or fixes across XLA frontends.

My guess is that (1) is the best approach to start with. On that front, I just noticed that @fehiepsi added some calculations involving incomplete gamma expressions in random.py, so that might be a good source of clues.

For approach (2) this seems like the relevant part of the XLA client math library but a quick glance didn’t turn up incomplete gamma stuff. Do you see any?

So overall my recommendation is to try implementing an incomplete gamma function in terms of lax.py primitives!

[hawkinsp]

This was fixed by @srvasude in PR #2064 .