Add support for 64-bit FFTs. (google#3290)

WORKSPACE

@@ -28,10 +28,10 @@ http_archive(
 #    and update the sha256 with the result.
 http_archive(
     name = "org_tensorflow",
-    sha256 = "642f5a1bc191dfb96b2d7ed1cfb8f2a1515b5169b8de4381c75193cef8404b92",
-    strip_prefix = "tensorflow-b25fb1fe32094b60f5a53ad5f986ad65a9f05919",
+    sha256 = "99231c027ad22e1a82866d2e6bc60379d06d0a75793ac09b547282eb5b382d37",
+    strip_prefix = "tensorflow-37aaafb0c1baa7acd0607748326cc12faf556277",
     urls = [
-        "https://github.com/tensorflow/tensorflow/archive/b25fb1fe32094b60f5a53ad5f986ad65a9f05919.tar.gz",
+        "https://github.com/tensorflow/tensorflow/archive/37aaafb0c1baa7acd0607748326cc12faf556277.tar.gz",
     ],
 )
 

jax/lax/lax_fft.py

@@ -23,6 +23,7 @@
 from jax.interpreters import xla
 from jax.util import prod
 from . import dtypes, lax
+from .. import lib
 from ..lib import xla_client
 from ..interpreters import ad
 from ..interpreters import batching
@@ -35,16 +36,18 @@
 ]
 
 def _promote_to_complex(arg):
-  dtype = onp.result_type(arg, onp.complex64)
-  # XLA's FFT op only supports C64.
-  if dtype == onp.complex128:
+  dtype = dtypes.result_type(arg, onp.complex64)
+  # XLA's FFT op only supports C64 in jaxlib versions 0.1.47 and earlier.
+  # TODO(phawkins): remove when minimum jaxlib version is 0.1.48 or newer.
+  if lib.version <= (0, 1, 47) and dtype == onp.complex128:
     dtype = onp.complex64
   return lax.convert_element_type(arg, dtype)
 
 def _promote_to_real(arg):
-  dtype = onp.result_type(arg, onp.float64)
+  dtype = dtypes.result_type(arg, onp.float64)
   # XLA's FFT op only supports F32.
-  if dtype == onp.float64:
+  # TODO(phawkins): remove when minimum jaxlib version is 0.1.48 or newer.
+  if lib.version <= (0, 1, 47) and dtype == onp.float64:
     dtype = onp.float32
   return lax.convert_element_type(arg, dtype)
 

jaxlib/version.py

@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "0.1.47"
+__version__ = "0.1.48"

tests/fft_test.py

@@ -20,6 +20,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 
+import jax
 from jax import lax
 from jax import numpy as jnp
 from jax import test_util as jtu
@@ -29,9 +30,7 @@
 
 
 float_dtypes = [np.float32, np.float64]
-# TODO(b/144573940): np.complex128 isn't supported by XLA, and the JAX
-# implementation casts to complex64.
-complex_dtypes = [np.complex64]
+complex_dtypes = [np.complex64, np.complex128]
 inexact_dtypes = float_dtypes + complex_dtypes
 int_dtypes = [np.int32, np.int64]
 bool_dtypes = [np.bool_]