Open notlober opened 1 month ago
@will-cromar should be able to share some information.
All three frameworks interact with libtpu
through the PJRT plugin API. Most of the core API for PJRT is documented in comments here: https://github.com/openxla/xla/blob/main/xla/pjrt/pjrt_client.h
Almost all of our interactions with PJRT are in this folder, and it's largely independent from PyTorch itself: https://github.com/pytorch/xla/tree/master/torch_xla/csrc/runtime
Specifically, to create a PJRT TPU client, you would need to go through the PjRtCApiClient
similar to this (device_type = "tpu"
, library_path = "/path/to/libtpu.so"
): https://github.com/pytorch/xla/blob/dd3b00ca95f455ffd9eec15803429420efbd106a/torch_xla/csrc/runtime/pjrt_registry.cc#L118-L126
Once you have a client instantiated, then your interactions are going to look a lot like this example from JAX: https://github.com/google/jax/blob/main/examples/jax_cpp/main.cc
We use the PJRT C++ API direcly, but it's worth noting that (other than the example above) JAX actually mainly interacts with PJRT through Python bindings. I not nearly as familiar with those, so you'll have better luck asking in their repository if you want to use the same bindings.
The framework code outside of libtpu.so
is all open source. I'm happy to help if you have any questions about the PJRT C++ API.
I'm looking to build an automatic differentiation library for TPUs without using high-level front-ends like TensorFlow/JAX/PyTorch-XLA, but I'm finding information about lower-level TPU usage is practically non-existent.
Specifically, I'm interested in:
Are there any insights or suggestions on how to approach this, particularly regarding TPU support? Any ideas or help would be greatly appreciated.
I understand that some of this information might be proprietary, but any guidance on what is possible or available would be very helpful.