[fritzo]

PyTorch has had an XLA backend for years. I don't know how performant it is though. https://pytorch.org/xla

[m00x]

It's pretty fast, just not as nice to use. You need statically defined tensors, and some functions are just not supported (last time I used it).