Fix log1p inaccuracies on complex inputs with large absolute values.

[pearu]

As in the title.

Tests and improvement reports are in google/jax#20144.

Accuracy tests are enabled in google/jax#20436

[pearu]

The same comment as in #10376 (comment) applies here as well: the CI testOnComplexPlane failures are expected and that should be fixed after google/jax#20144 lands

[cheshire]

Could we test the improved accuracy?

[pearu]

Could we test the improved accuracy?

The accuracy of log1p is tested in google/jax#20144 . See the description of that PR that reports the accuracy improvements for log1p.

[tdanyluk]

@cheshire gentle reminder

[cheshire]

The accuracy of log1p is tested in google/jax#20144 . See the description of that PR that reports the accuracy improvements for log1p.

Understood, so JAX test is an e2e integration test, but could we have an XLA unit test as well?

[pearu]

... could we have an XLA unit test as well?

@cheshire Having complex functions correctness tests in XLA rather than in JAX makes lots of sense.

That said, when testing functions for correctness, it is essential to have reference implementations that correctness are verified in a platform independent way. Unfortunately, there are not many libraries available that qualify as a reference library.

For instance, within the Python/JAX framework, one obvious choice would be numpy, however, as it turns out, numpy functions often rely on system math libraries and the function evaluation results can be different on different platforms. This is one important reason why in JAX tests we use mpmath as a reference library. While there exists a few bugs in mpmath itself that require workarounds, one of the most important properties of mpmath is that its functions evaluation results are platform independent. In https://github.com/pearu/complex_function_validation/, I have validated complex functions from a number of Python array libraries for correctness with a conclusion that mpmath is a reasonable reference library indeed.

When considering XLA, one of the obvious choices is libc++. Unfortunately, there exists incorrectness issues also in libc++ that must be tackled via correctness analysis and implementing workarounds before using it as a reference library. The current tests in XLA are very minimal (functions are tested on a couple of points only) and miss many incorrectness issues that the correctness analysis of JAX complex functions above also demonstrates.

In short term, one approach would be to generate tables of inputs and outputs to different complex functions (say, using mpmath) and use these tables as a reference database when testing XLA complex functions for correctness.

What do you think?

[cheshire]

In this case having hardcoded expected output values along with the long comment as the one above indeed seems like the preferred option, thanks!

[pearu]

In this case having hardcoded expected output values along with the long comment as the one above indeed seems like the preferred option, thanks!

I suggest to do this in a follow up PR because it is relevant for other functions as well, and to avoid XLA and JAX repos going out of sync (google/jax#20436 has landed that includes xla_extension_version update related to this PR).

[cheshire]

Sorry what's the connection between going out of sync and having a unit test?

[pearu]

Sorry what's the connection between going out of sync and having a unit test?

In this particular case, google/jax#20436 enabled unit tests for log1p fixes that are implemented in this PR. Atm, this means that when running JAX tests from its main branch using jaxlib built from XLA main branch, the JAX log1p unit tests will fail. xla_extension_version is designed to prevent this but under assumption that the related XLA and JAX PRs are landed more or less simultaneously. When some unrelated XLA PR updates xla_extension_version and it lands before, then the related XLA and JAX PRs will be out of sync. The current situation is a bit worse: XLA and JAX main branches are out of sync.

[pearu]

The current situation is a bit worse: XLA and JAX main branches are out of sync.

A quick solution to this is to create a JAX PR that disables log1p units tests. Although, it would also make sense to land this PR as we already do test the related changes in JAX. And in the future, we'll implement the corresponding unit tests also in XLA. What do you think?

[cheshire]

Sorry I'm still confused. This PR improves log1p handling in XLA, right? So it would be natural to accompany it with a unit test? The question regarding JAX and/or JAX/XLA interop is orthogonal?

[pearu]

Sorry I'm still confused. This PR improves log1p handling in XLA, right?

Correct.

So it would be natural to accompany it with a unit test?

True. I am working on it atm.

The question regarding JAX and/or JAX/XLA interop is orthogonal?

There exists inference only when JAX implements tests for XLA features. AFAIK, this situation is not unique to this issue.