[jax2tf] First draft of testing the QR conversion. #3775
Conversation
Complex QR would not work if compiled with XLA in TensorFlow because there is no XLA-level implementation. It works in JAX on CPU and GPU because we call LAPACK and Cusolver via custom-calls. It does not work in JAX on TPU for the same reason (#1274): we don't have an implementation yet. Also note that even for real QR decompositions you will call the HLO implementation when compiling with TF, which is probably much slower than the LAPACK/Cusolver implementations. Arguably the right fix for that is to upstream our linear algebra implementations to XLA so TF/XLA can call them also. |
| self.ConvertAndCompare(harness.dyn_fun, *harness.dyn_args_maker(self.rng()), | ||
| atol=1e-5, rtol=1e-5) | ||
| else: | ||
| # TODO: see https://github.com/google/jax/pull/3775#issuecomment-659407824. |
There was a problem hiding this comment.
This is a good strategy to keep track of TODOs that are discussed in comments.
QR decomposition is off by over 1e-6 in some instances, requiring custom atol and rtol values in testing code. There is an odd problem in which experimental compilation fails for complex types, although they are in principle supported.
|
I tested this internally before merging. |
* [jax2tf] First draft of testing the QR conversion. QR decomposition is off by over 1e-6 in some instances, requiring custom atol and rtol values in testing code. There is an odd problem in which experimental compilation fails for complex types, although they are in principle supported.
* [jax2tf] First draft of testing the QR conversion. QR decomposition is off by over 1e-6 in some instances, requiring custom atol and rtol values in testing code. There is an odd problem in which experimental compilation fails for complex types, although they are in principle supported.
* [jax2tf] First draft of testing the QR conversion. QR decomposition is off by over 1e-6 in some instances, requiring custom atol and rtol values in testing code. There is an odd problem in which experimental compilation fails for complex types, although they are in principle supported.
QR decomposition is off by over 1e-6 in some instances, requiring custom
atol and rtol values in testing code. There is an odd problem in which
experimental compilation fails for complex types, although they are in
principle supported.