How to use a PJRT plugin from Pytorch?

[wwtghx]

In the file https://github.com/openxla/xla/blob/main/xla/pjrt/c/docs/pjrt_integration_guide.md, "How to use a PJRT plugin from JAX" is found, but I want to know "How to use a PJRT plugin from Pytorch"?

[cheshire]

I think the expected usage pattern is via https://github.com/pytorch/xla, which itself should use the PJRT API.

[cheshire]

@JackCaoG for further comments

[JackCaoG]

Depending on what you mean by use. If you are a user who wants to use PyTorch/XLA with PJRT, it is already the default runtime after 2.1 release. Check out our doc at https://github.com/pytorch/xla/blob/master/docs/pjrt.md#pjrt-runtime.

If you are a backend vender who want to integrate PyTorch/XLA with your backend. @will-cromar recently started the effort of developing the python PJRT plugin for PyTorch/XLA. My understanding is that registration process will be different, but you would also need to implement the PJRT CAPI similar to https://github.com/openxla/xla/blob/main/xla/pjrt/c/docs/pjrt_integration_guide.md#step-1-implement-pjrt-c-api-interface .

[wwtghx]

Yes, I am a backend vender, and I want to integrate PyTorch/XLA with my backend.

[JackCaoG]

Great, we have helped a couple backend vender to integrated into pytorch/xla. @will-cromar recently started the python plugin effort in pytorch/xla#5644, I expect it to be more mature in a couple months. In the mean time I think you can start with implement the PJRT CAPIs, after that we can help with the backend registration logic.

[will-cromar]

Hey, @wwtghx. It's possible today to use a PJRT Plugin from PyTorch/XLA, but the process is rough around the edges. It basically amounts to updating lists and special cases throughout this repository as in pytorch/xla#5428.

This is an active area of work, and the end goal is to simplify the process and implement a registration mechanism for external plugins like JAX has. I started pytorch/xla#6242 to collect my ideas and track progress. Please take a look and let me know if you have any feedback on the design.