Is there float32 precision on Google TPU?

Question

I plan to use Google TPUs for scientific numerical simulation (finite element analysis).

That said, do TPUs support float32 for computation? If the precision is too small its unsuitable for me.

There is a notification here, but it says XLA converts it automatically.

Matthew Flaschen · Accepted Answer · 2024-09-15 16:27:14Z

0

Yes. It is not the default, but it is available ("Run a calculation on a Cloud TPU VM using JAX" at Google Cloud).

You can set precision=jax.lax.Precision.HIGHEST on particular operations, such as matrix multiplication.

It "uses even more [Matrix Multiply Units] passes to achieve full float32 precision".

E.g.:

jax.numpy.matmul(a, b, precision=Precision.HIGHEST)

answered Sep 14 at 20:10

Matthew Flaschen

284k53 gold badges520 silver badges552 bronze badges

Add a comment |

user27327621user27327621 · Accepted Answer · 2024-09-14 21:29:31Z

0

If you are able to do so on that platform, try checking if sizeof(float)*CHAR_BIT == 32

answered Sep 14 at 21:29

user27327621

TPU mathematical operations use a custom compiler, not standard C data types: "Code that runs on TPUs must be compiled by the accelerator linear algebra (XLA) compiler. XLA is a just-in-time compiler that takes the graph emitted by an ML framework application and compiles the linear algebra, loss, and gradient components of the graph into TPU machine code." (cloud.google.com/tpu/docs/intro-to-tpu#xla_compiler).
– Matthew Flaschen
Commented Sep 15 at 16:46

Add a comment |

2 Answers 2