Support erfcx

[cossio]

Jax is missing support for scipy.special.erfcx (https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.erfcx.html). Please add this function.

[cossio]

@mattjj I see this was marked a "good first issue". If you think this is easy to do, I could be interested in trying. Can you point me to docs or something I can read?

[hawkinsp]

Well, ultimately how these things are implemented is either:
a) via expressing them in terms of other preexisting mathematical functions (here exp and erfc), or
b) as some sort of polynomial expansion.

Here you certainly could do (a), but you might hit underflow and overflow problems, since presumably they are the reason the composed erfcx function exists in the first place.

Happily the following stack overflow question looks like it has some good pointers, including to a paper that describes how to implement erfcx in terms of Chebyshev polynomials:
https://stackoverflow.com/questions/39777360/accurate-computation-of-scaled-complementary-error-function-erfcx

To write such an implementation in JAX, you could simply express that algorithm in Python using lax and jax.numpy primitives.

Does that help?

[hawkinsp]

(Note you can't necessarily use the implementation in the StackOverflow answer verbatim, not least because JAX doesn't have a fused multiply-add operator.)

[cossio]

@hawkinsp Yes, thanks. I have some ideas on how to implement erfcx generically. My question was mostly about how to extend jax? Can you expand a bit on that? What docs should I read?

Maybe there are some PRs implementing other special functions that one can use as guiding examples

[hawkinsp]

Sure! Take a look at:
https://github.com/google/jax/blob/master/jax/scipy/special.py
which has an implementation of, say, ndtri. It really is as simple as defining a Python function in that file.

Note also that JAX actually gets many of these from XLA's C++ client library. This is exactly the same idea, but implemented in C++ (and included in jaxlib).
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/client/lib/math.h
The relative performance of C++ vs Python doesn't matter in JAX if you wrap the implementation in jax.jit since all we are doing is building a compiler IR; you can send a PR wherever you think is more convenient in whichever language you are more comfortable. I can walk you through the XLA plumbing if you prefer to do it in C++.

[srvasude]

FWIW, I should shortly be contribution an erfcx implementation in to TensorFlow Probability (to fix some numerics in our Normal-adjacent distributions) which should be accessible in our JAX backend.

[jonas-eschle]

FWIW, I should shortly be contribution an erfcx implementation in to TensorFlow Probability (to fix some numerics in our Normal-adjacent distributions) which should be accessible in our JAX backend.

That would be great! However, will it support complex arguments as well? So far, erfc is unfortunately missing support for them.

That would be quite useful in order to have the faddeeva function available.

[srvasude]

At the moment it would not, since our general target has been supporting float/double arguments (since these are generally being used in Normal-adjacent distributions).

[aritropc]

Is any task left for this issue? If yes, I'd like to work on that

[jonas-eschle]

Hi @aritropc there is indeed some work left: We explored the possibility to add it to TensorFlow Probability (which has a JAX substrate, so it would be possible to use it in jax as well) in a GSoC project that is located here: https://github.com/Aman027/GSOC-21-Report?ref=https://githubhelp.com

We're currently discussing the implementation into TFP here: tensorflow/probability#1265

It got a bti stuck however and it would be great if you could give it a go. Maybe just start writing in tensorflow/probability#1265