Google的TPU的Pallas扩展功能支持的数据类型

地址：

https://jax.readthedocs.io/en/latest/pallas/tpu.html

jnp.float32

jnp.bfloat16

jnp.int* (all precisions, except for jnp.int4)

jnp.uint* (all precisions)