Implement complex QR decomposition in HLO (TPU)

[mganahl]

The following code raises RuntimeError: Unimplemented: complex comparison 'LT'

import jax
jax.np.linalg.qr((np.random.rand(3, 3) + 1j * np.random.rand(3, 3)))

Noticed this because one of our test cases fails (google/TensorNetwork#221)

[hawkinsp]

The QR implementation we are using is implemented in XLA/HLO on all platforms and only supports real inputs.

On CPU and GPU we should really be calling LAPACK and Cusolver implementations of QR decomposition anyway. Both libraries support complex inputs and would probably perform better than the HLO implementation on their respective platforms.

On TPU we would need to add complex number support here:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/client/lib/qr.cc

Which platform are you interested in?

[mganahl]

Hi! We'd be interested in all three. The ranking is CPU, GPU and TPU, with TPU the least important right now. Thanks for the quick reply!

[hawkinsp]

PR #1306 adds complex QR support on CPU and GPU.

TPU is not yet implemented, so I'll leave this issue open and retitle it.

[mganahl]

awesome thanks!

[long1216]

I got a similar error, from #3862, RuntimeError: Unimplemented: complex comparison 'GE'
Does they belong to the same problem, also related to the complex comparison in XLA? What the meaning of the terms 'GE', 'LT' stand for? Thank you so much.

[mattjj]

GE and LT are "greater than" and "less than", respectively. That error is saying that those operations aren't implemented for complex numbers. (NumPy implements them with a kind of weird convention, since complex numbers don't have a total ordering.)

[long1216]

Thank you for your answer, Matthew. Would there be any plans for finalizing the conversion of complex comparison? Could we just follow the convention of NumPy in XLA first? Thank you.

[hawkinsp]

We've already fixed this at head!