Will 64-bit data type support be added to the TPU driver API?

[willfarris]

Hi,

When I try to create a jax array of type float64 in Google Colab with the TPU runtime selected I get the following message:

RuntimeError: INVALID_ARGUMENT: 64-bit data types are not yet supported on the TPU driver API. Convert inputs to float32/int32_t before using.

The verbiage here suggests this support may one day be added to the TPU driver API, so I'd like to know if this is something that's planned to be added in the future. If it is planned, is there an ETA on when this feature may be available? I opened an issue on the Tensorflow repo and was told this would be a more appropriate place to ask.

For a university project I'm trying to implement a Mersenne prime number search program which runs on the TPU. These calculations are very sensitive to precision errors, and we've found that although we may be able to use the TPU to improve performance over existing GPU-based programs, the precision of bfloat16 or even float32 are too low for the calculations we're trying to run.

running the code below in a Google Colab notebook with the TPU runtime selected should produce the error:

import jax
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()
jax.config.update("jax_enable_x64", True)
x = jax.numpy.array([1, 2, 3, 4], dtype=jax.numpy.float64) # Should throw the exception mentioned above