chore: import upstream snapshot with attribution

third_party/xla/docs/pjrt/_toc.yaml

@@ -0,0 +1,27 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+toc:
+- heading: PJRT
+- title: Getting started
+  section:
+  - title: Introduction
+    path: /xla/pjrt
+  - title: PJRT C++ API Overview
+    path: /xla/pjrt/cpp_api_overview
+  - title: Develop a New PJRT Plugin
+    path: /xla/pjrt/pjrt_integration
+  - title: PJRT Examples
+    path: /xla/pjrt/examples

third_party/xla/docs/pjrt/cpp_api_overview.md

@@ -0,0 +1,333 @@
+# PJRT C++ Device API Overview
+
+## Background
+
+[PJRT](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h)
+is the uniform Device API that we want to add to the ML ecosystem. The long term
+vision is that:
+
+1. Frameworks (JAX, TF, etc.) will call PJRT, which has device-specific
+   implementations that are opaque to the frameworks;
+2. Each device focuses on implementing PJRT APIs, and can be opaque to the
+   frameworks.
+
+PJRT offers both a C API and C++ API. Plugging in at either layer is OK, the C++
+API uses classes to abstract away some concepts, but also has stronger ties to
+XLA datatypes. This page focuses on the C++ API.
+
+## PJRT Components
+
+![PJRT Components](images/pjrt_client.svg)
+
+Note: Most items in this diagram also have backpointers, memory spaces know
+their device(s) and client, devices know their client, buffers know their memory
+space.
+
+### PjRtClient
+
+_Full reference at [`pjrt_client.h > PjRtClient`](https://github.com/openxla/xla/blob/924b74d84de3760cc589fd1525c7346691d51df5/xla/pjrt/pjrt_client.h#L486)._
+
+Clients manage all communication between the device and framework, and
+encapsulate all state used in the communication. They have a generic set of APIs
+for interacting with a PJRT plugin, and they own the devices and memory spaces
+for a given plugin.
+
+### PjRtDevice
+
+_Full references at [`pjrt_client.h > PjRtDevice`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_client.h#L102),
+and [`pjrt_device_description.h`](https://github.com/openxla/xla/blob/main/xla/pjrt/pjrt_device_description.h)_
+
+A device class is used to describe a single device. A device has a device
+description to help identify its kind (unique hash to identify GPU/CPU/xPU), and
+location within a grid of devices both locally and globally.
+
+Devices also know their associated memory spaces and the client it is owned by.
+
+A device does *not* necessarily know the buffers of actual data associated with
+it, but it can figure that out by looking through its associated memory spaces.
+
+### PjRtMemorySpace
+
+_Full reference at [`pjrt_client.h > PjRtMemorySpace`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_client.h#L72)._
+
+Memory spaces can be used to describe a location of memory. These can either be
+unpinned, and are free to live anywhere but be accessible from a device, or they
+can be pinned and must live on a specific device.
+
+Memory spaces know their associated buffers of data, and the devices (plural)
+that a memory space is associated with, as well as the client it is a part of.
+
+### PjRtBuffer
+
+_Full reference at [`pjrt_client.h > PjRtBuffer`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_client.h#L1111)._
+
+A buffer holds data on a device in some format that will be easy to work with
+inside the plugin, such as an MLIR elements attr or a proprietary tensor format.
+A framework may try to send data to a device in the form of an `xla::Literal`,
+i.e. for an input argument to the module, which must be cloned (or borrowed), to
+the device's memory. Once a buffer is no longer needed the `Delete` method is
+invoked by the framework to clean up.
+
+A buffer knows the memory space it is a part of, and transitively can figure out
+which devices are able to access it, but buffers don't necessarily know their
+devices.
+
+For communicating with frameworks, buffers know how to convert to and from an
+`xla::Literal` type:
+
+```cpp
+// Literal to Buffer
+absl::StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(...) {...}
+
+// Buffer to Literal
+xla::Future<> ToLiteral(xla::MutableLiteralBase* literal) override {...}
+```
+
+APIs for creating a buffer have [Buffer Semantics](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_client.h#L858)
+which help dictate if literal data from the host buffer can be shared or copied
+or mutated.
+
+Lastly, a buffer may need last longer than the scope of its execution, if it is
+assigned to a variable in the framework layer `x = jit(foo)(10)`, in these cases
+buffers allow building external references which provide a temporarily owned
+pointer to the data held by the buffer, along with metadata (dtype / dim sizes)
+for interpreting the underlying data.
+
+### PjRtCompiler
+
+_Full reference at [`pjrt_compiler.h > PjRtCompiler`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_compiler.h#L157)._
+
+The `PjRtCompiler` class provides useful implementation details for XLA
+backends, but is not necessary for a plugin to implement. In theory, the
+responsibility of a `PjRtCompiler`, or the `PjRtClient::Compile` method, is to
+take an input module and return a `PjRtLoadedExecutable`.
+
+### PjRtExecutable / PjRtLoadedExecutable
+
+_Full reference at [`pjrt_executable.h > PjRtExecutable`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_executable.h#L306),
+and [`pjrt_client.h > PjRtLoadedExecutable`](https://github.com/openxla/xla/blob/3e448cf9e86775a37ec5f7d3c69dfb20e0c760df/xla/pjrt/pjrt_client.h#L1506)._
+
+A `PjRtExecutable` knows how to take a compiled artifact and execution options
+and serialize/deserialize them so an executable can be stored and loaded as
+needed.
+
+The `PjRtLoadedExecutable` is the in-memory compiled executable which is ready
+for input arguments to execute, it is a subclass of `PjRtExecutable`.
+
+Executables are interfaced with via one of the client's `Execute` methods:
+
+```cpp
+// Execute on addressable devices
+absl::StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
+Execute(absl::Span<const std::vector<PjRtBuffer*>> argument_handles, ...) {...}
+
+// Execute assigned replica/partition on the specified device
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+ExecuteSharded(absl::Span<PjRtBuffer* const> argument_handles,
+              PjRtDevice* device, ...) {...}
+
+// Execute on specified device, single replica / partition
+absl::StatusOr<std::vector<std::unique_ptr<PjRtBuffer>>>
+ExecutePortable(absl::Span<PjRtBuffer* const> argument_handles,
+                PjRtDevice* device, ...) {...}
+```
+
+Before calling `Execute` the framework will transfer all required data to
+`PjRtBuffers` owned by the executing client, but returned for the framework to
+reference. These buffers are then provided as arguments to the `Execute` method.
+
+## PJRT Concepts
+
+### Futures & Async Computations
+
+If any part of a plugin is implemented asynchronously, it _must_ properly
+implement futures.
+
+Consider the following program:
+
+```py
+@jax.jit
+def foo(x): return x + 1
+
+x = foo(1)
+# [...] other logic not using `x`
+print(x + 1)
+```
+
+An async plugin would be able to enqueue the computation `x`, and immediately
+return a buffer which isn't ready to be read yet, but execution will populate
+it. Execution can continue to enqueue necessary computations after `x`, that
+don't require `x`, including execution on other PJRT devices. Once the value of
+`x` is needed, execution will block until the buffer declares itself ready via
+the future returned by `GetReadyFuture`.
+
+Futures can be useful to determine when an object becomes available, including
+devices and buffers.
+
+### Advanced concepts
+
+Extending beyond implementing the base APIs will expand the features of JAX that
+can be used by a plugin. These are all opt-in features in the sense that at
+typical JIT and execute workflow will work without them, but for a production
+quality pipeline some thought should likely be put into the degree of support
+for any of these features supported by PJRT APIs:
+
+- Memory spaces
+- Custom layouts
+- Communication ops like send/recv
+- Host offloading
+- Sharding
+
+## Typical PJRT framework-device communication
+
+### Example Log
+
+The following is a log of the methods called to load the PJRT plugin and
+execute `y = jax.jit(lambda x: jnp.power(x, jnp.int32(2)))(1)`. In this case
+we log JAX interacting with the StableHLO Reference PJRT plugin.
+
+<details>
+<summary>Example log</summary>
+<br>
+<pre>
+
+```
+//////////////////////////////////
+// Load the plugin
+//////////////////////////////////
+
+I client_cpp_pjrt.cc:55] StablehloReferencePjrtClient(0x23bac400)
+I device.cc:53] StablehloReferenceDeviceDescription(0x23bac4f8)
+I device.cc:104] StablehloReferenceDevice(0x23bac4e0)
+I device.cc:123] client(0x23bac4e0)
+I device.cc:123] client(0x23bac4e0)
+I client_cpp_pjrt.cc:71] process_index(0x23bac400)
+I client_cpp_pjrt.cc:67] platform_name(0x23bac400)
+I device.cc:143] AttachDefaultMemorySpace(0x23bac4e0)
+I client_cpp_pjrt.cc:67] platform_name(0x23bac400)
+I client_cpp_pjrt.cc:86] devices(0x23bac400)
+I client_cpp_pjrt.cc:81] addressable_device_count(0x23bac400)
+I device.cc:168] description(0x23bac4e0)
+I device.cc:168] description(0x23bac4e0)
+I device.cc:86] Attributes(0x23bac4f8)
+I device.cc:128] IsAddressable(0x23bac4e0)
+I device.cc:168] description(0x23bac4e0)
+I device.cc:61] process_index(0x23bac4f8)
+I device.cc:123] client(0x23bac4e0)
+I client_cpp_pjrt.cc:71] process_index(0x23bac400)
+I client_cpp_pjrt.cc:81] addressable_device_count(0x23bac400)
+I client_cpp_pjrt.cc:95] memory_spaces(0x23bac400)
+I device.cc:128] IsAddressable(0x23bac4e0)
+I device.cc:168] description(0x23bac4e0)
+I device.cc:61] process_index(0x23bac4f8)
+I device.cc:123] client(0x23bac4e0)
+I client_cpp_pjrt.cc:71] process_index(0x23bac400)
+I device.cc:148] memory_spaces(0x23bac4e0)
+Creating PJRT Client from client
+I client_cpp_pjrt.cc:108] platform_version(0x23bac400)
+I client_cpp_pjrt.cc:67] platform_name(0x23bac400)
+I device.cc:57] id(0x23bac4f8)
+I device.cc:70] device_kind(0x23bac4f8)
+I device.cc:70] device_kind(0x23bac4f8)
+I device.cc:80] ToString(0x23bac4f8)
+I device.cc:80] ToString(0x23bac4f8)
+I device.cc:75] DebugString(0x23bac4f8)
+I device.cc:75] DebugString(0x23bac4f8)
+I device.cc:61] process_index(0x23bac4f8)
+I device.cc:128] IsAddressable(0x23bac4e0)
+I device.cc:168] description(0x23bac4e0)
+I device.cc:61] process_index(0x23bac4f8)
+I device.cc:123] client(0x23bac4e0)
+I client_cpp_pjrt.cc:71] process_index(0x23bac400)
+I device.cc:153] default_memory_space(0x23bac4e0)
+I client_cpp_pjrt.cc:71] process_index(0x23bac400)
+
+//////////////////////////////////
+// RUN: `y = jax.jit(lambda x: jnp.power(x, jnp.int32(2)))(1)`
+//////////////////////////////////
+
+I executable.cc:309] num_partitions(0x240bab70)
+I executable.cc:305] num_replicas(0x240bab70)
+I executable.cc:309] num_partitions(0x240bab70)
+I client_cpp_pjrt.cc:233] BufferFromHostBuffer(0x23bac400)
+I buffer.cc:285] CreateMlirBufferFromLiteral
+I buffer.cc:98] CreateFromLiteral
+I buffer.cc:99] CreateFromLiteral: s32[] 2
+I buffer.cc:64] MlirPjrtBuffer(0x240bb050)
+I buffer.cc:102] CreateFromLiteral -> 0x240bb050
+I buffer.cc:158] device(0x240bb050)
+I buffer.cc:154] memory_space(0x240bb050)
+I buffer.cc:154] memory_space(0x240bb050)
+I executable.cc:328] GetHloModules(0x240bab70)
+I executable.cc:240] Execute(0x240bab70)
+I executable.cc:197] ExecuteWithReferenceInterpreter(0x240bab70)
+I buffer.cc:303] GetAttributeFromBuffer
+I buffer.cc:229] IsDeleted(0x240bb050)
+I buffer.cc:311] GetAttributeFromBuffer(0x240bb050) -> dense<2> : tensor<i32>
+I executable.cc:205] EvalModule:
+module @jit attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
+  func.func public @main(%arg0: tensor<i32> {mhlo.layout_mode = "default"}) -> (tensor<i32> {jax.result_info = "", mhlo.layout_mode = "default"}) {
+    // ...
+    return %3 : tensor<i32>
+  }
+}
+I executable.cc:206] Inputs: [dense<2> : tensor<i32>]
+I executable.cc:213] Results: [dense<2> : tensor<i32>]
+I device.cc:153] default_memory_space(0x23bac4e0)
+I buffer.cc:291] CreateMlirBufferFromAttribute
+I buffer.cc:116] CreateFromAttribute
+I buffer.cc:64] MlirPjrtBuffer(0x22cea630)
+I buffer.cc:122] CreateFromAttribute(dense<2> : tensor<i32>) -> 0x22cea630
+
+//////////////////////////////////
+// RUN: `print(y)`
+//////////////////////////////////
+
+I buffer.cc:263] GetReadyFuture(0x22cea630)
+I buffer.cc:264] GetReadyFuture(0x22cea630)
+I buffer.cc:154] memory_space(0x22cea630)
+I buffer.cc:154] memory_space(0x22cea630)
+I buffer.cc:158] device(0x22cea630)
+I buffer.cc:158] device(0x22cea630)
+I buffer.cc:154] memory_space(0x22cea630)
+I buffer.cc:154] memory_space(0x22cea630)
+I buffer.cc:229] IsDeleted(0x22cea630)
+I buffer.cc:129] on_device_shape(0x22cea630)
+I buffer.cc:129] on_device_shape(0x22cea630)
+I buffer.cc:129] on_device_shape(0x22cea630)
+I buffer.cc:158] device(0x22cea630)
+I buffer.cc:154] memory_space(0x22cea630)
+I buffer.cc:154] memory_space(0x22cea630)
+I client_cpp_pjrt.cc:71] process_index(0x23bac400)
+I buffer.cc:229] IsDeleted(0x22cea630)
+I buffer.cc:129] on_device_shape(0x22cea630)
+I buffer.cc:129] on_device_shape(0x22cea630)
+I buffer.cc:269] IsOnCpu(0x22cea630) # Returns true, allows external references.
+I buffer.cc:129] on_device_shape(0x22cea630)
+I buffer.cc:129] on_device_shape(0x22cea630)
+I buffer.cc:129] on_device_shape(0x22cea630)
+I buffer.cc:129] on_device_shape(0x22cea630)
+I buffer.cc:129] on_device_shape(0x22cea630)
+I buffer.cc:168] AcquireExternalReference(0x22cea630)
+I buffer.cc:73] MlirClonedExternalReference(0x2404d560)
+I buffer.cc:303] GetAttributeFromBuffer
+I buffer.cc:229] IsDeleted(0x22cea630)
+I buffer.cc:311] GetAttributeFromBuffer(0x22cea630) -> dense<2> : tensor<i32>
+I buffer.cc:291] CreateMlirBufferFromAttribute
+I buffer.cc:116] CreateFromAttribute
+I buffer.cc:64] MlirPjrtBuffer(0x240bb050)
+I buffer.cc:122] CreateFromAttribute(dense<2> : tensor<i32>) -> 0x240bb050
+I buffer.cc:168] AcquireExternalReference(0x22cea630)
+I buffer.cc:73] MlirClonedExternalReference(0x240b6010)
+I buffer.cc:303] GetAttributeFromBuffer
+I buffer.cc:229] IsDeleted(0x22cea630)
+I buffer.cc:311] GetAttributeFromBuffer(0x22cea630) -> dense<2> : tensor<i32>
+I buffer.cc:291] CreateMlirBufferFromAttribute
+I buffer.cc:116] CreateFromAttribute
+I buffer.cc:64] MlirPjrtBuffer(0x23b2db60)
+I buffer.cc:122] CreateFromAttribute(dense<2> : tensor<i32>) -> 0x23b2db60
+I buffer.cc:263] GetReadyFuture(0x22cea630)
+I buffer.cc:264] GetReadyFuture(0x22cea630)
+```
+</pre>
+</details>

third_party/xla/docs/pjrt/examples.md

@@ -0,0 +1,38 @@
+# PJRT Examples
+
+## Example: JAX CUDA plugin
+
+1. PJRT C API implementation through wrapper ([pjrt\_c\_api\_gpu.h](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api_gpu.h)).
+1. Set up the entry point for the package ([setup.py](https://github.com/google/jax/blob/main/jax_plugins/cuda/setup.py)).
+1. Implement an initialize() method ([\_\_init\_\_.py](https://github.com/google/jax/blob/a10854786b6d1bc92a65dd314916b151640789af/plugins/cuda/__init__.py#L31-L51)).
+1. Can be tested with any jax tests for CUDA.
+
+
+## Frameworks Implementations
+
+Some references for using PJRT on the framework side, to interface with PJRT
+devices:
+
+- JAX
+  + [jax-ml/jax](https://github.com/jax-ml/jax/blob/main/jax/_src/compiler.py#L248)
+    interacts with PJRT APIs via the `xla_client` APIs
+- GoMLX
+  + [gomlx/gopjrt](https://github.com/gomlx/gopjrt)
+  + [gomlx/gomlx > backends/xla](https://github.com/gomlx/gomlx/tree/main/backends/xla/xla.go)
+- ZML
+  + PJRT API wrapper [pjrt.zig](https://github.com/zml/zml/blob/master/pjrt/pjrt.zig)
+  + Load PJRT Plugin [context.zig](https://github.com/zml/zml/blob/master/zml/context.zig#L30-L34)
+  + Interacting with PJRT Buffers [buffer.zig](https://github.com/zml/zml/blob/master/zml/buffer.zig#L36)
+  + Execute a module via PJRT [module.zig](https://github.com/zml/zml/blob/master/zml/module.zig#L863-L886)
+
+## Hardware Implementations
+
+- Full integration plugins (PJRT+MLIR+XLA):
+  + [XLA CPU Plugin](https://github.com/openxla/xla/tree/main/xla/pjrt/cpu/cpu_client.cc)
+  + [XLA GPU Plugin](https://github.com/openxla/xla/tree/main/xla/pjrt/gpu/se_gpu_pjrt_client.cc)
+  + [Intel XLA Plugin](https://github.com/intel/intel-extension-for-openxla)
+- Light integration plugins (PJRT+MLIR):
+  + StableHLO Reference Interpreter plugin
+    (MLIR-based, C++ plugin, to be linked after devlabs)
+  + [Tenstorrent-XLA plugin](https://github.com/tenstorrent/tt-xla/blob/main/src/common/pjrt_implementation/api_bindings.cc)
+    (MLIR-based, C plugin)

third_party/xla/docs/pjrt/index.md

@@ -0,0 +1,31 @@
+Project: /_project.yaml
+Book: /_book.yaml
+
+<link rel="stylesheet" href="/site-assets/css/style.css">
+
+# PJRT - Uniform Device API
+
+PJRT C API is the uniform Device API that we want to add to the ML ecosystem.
+The long term vision is that: (1) frameworks (TF, JAX, etc.) will call PJRT,
+which has device-specific implementations that are opaque to the frameworks; (2)
+each device focuses on implementing PJRT APIs as PJRT plugins, which can be
+opaque to the frameworks.
+
+## Communication channels
+
+* Issues and feature requests can be filed in the [OpenXLA/xla repo](https://github.com/openxla/xla).
+* Questions regarding PJRT can be asked on the [OpenXLA Discord][discord].
+
+## Resources
+
+*   [PJRT C API header](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h)
+*   [PJRT C API changelog](https://github.com/openxla/xla/blob/main/xla/pjrt/c/CHANGELOG.md)
+*   [PJRT integration guide](https://github.com/openxla/xla/blob/main/xla/pjrt/c/docs/pjrt_integration_guide.md)
+*   [PJRT design docs](https://drive.google.com/drive/folders/18M944-QQPk1E34qRyIjkqDRDnpMa3miN)
+*   [PJRT API ABI versioning and compatibility](https://docs.google.com/document/d/1TKB5NyGtdzrpgw5mpyFjVAhJjpSNdF31T6pjPl_UT2o/edit)
+*   [PJRT Plugin Mechanism design doc](https://docs.google.com/document/d/1Qdptisz1tUPGn1qFAVgCV2omnfjN01zoQPwKLdlizas/edit)
+*   [OpenXLA 2024 Fall DevLab PJRT plugin tutorial slides](https://drive.google.com/file/d/1epUJkMONG2t06WOeMHz4Oi3F_-8cTuz-/view)
+    ([recording](https://www.youtube.com/watch?v=2GlMqaNxP_w&list=PLlFotmaRrOzv2OIEpijqiHGmY7rpscFcj))
+*   [IREE PJRT plugin implementation](https://github.com/iree-org/iree/tree/main/integrations/pjrt)
+
+[discord]: https://discord.gg/ZKXq7b3V8A "Join on Discord"

third_party/xla/docs/pjrt/pjrt_integration.md

@@ -0,0 +1,153 @@
+# PJRT plugin integration
+
+This doc focuses on the recommendations about how to integrate with PJRT, and
+how to test PJRT integration with JAX.
+
+## How to integrate with PJRT
+
+### Step 1: Implement [PJRT C API interface](https://github.com/openxla/xla/blob/71a4e6e6e4e9f0f8b8f25c07a32ad489aff19239/xla/pjrt/c/pjrt_c_api.h)
+
+**Option A**: You can implement the PJRT C API directly.
+
+**Option B**: If you're able to build against C++ code in the [xla repo](https://github.com/openxla/xla) (via forking or bazel), you can also implement the PJRT C++ API and use the C→C++ wrapper:
+
+1. Implement a C++ PJRT client inheriting from the [base PJRT client](https://github.com/openxla/xla/blob/main/xla/pjrt/pjrt_client.h) (and related PJRT classes). Here are some examples of C++ PJRT client: [pjrt\_stream\_executor\_client.h](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/pjrt_stream_executor_client.h), [tfrt\_cpu\_pjrt\_client.h](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/tfrt_cpu_pjrt_client.h).
+1. Implement a few C API methods that are not part of C++ PJRT client:
+  * [PJRT\_Client\_Create](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h#L344-L365). Below is some sample pseudo code (assuming `GetPluginPjRtClient` returns a C++ PJRT client implemented above):
+
+```
+#include "third_party/tensorflow/compiler/xla/pjrt/c/pjrt_c_api_wrapper_impl.h"
+
+namespace my_plugin {
+PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args) {
+  std::unique_ptr<xla::PjRtClient> client = GetPluginPjRtClient();
+  args->client = pjrt::CreateWrapperClient(std::move(client));
+  return nullptr;
+}
+}  // namespace my_plugin
+```
+  Note [PJRT\_Client\_Create](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h#L344-L365) can take options passed from the framework. [Here](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api_gpu_internal.cc#L48-L102) is an example of how a GPU client uses this feature.
+
+  * [Optional] [PJRT\_TopologyDescription\_Create](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h#L1815-L1830).
+  * [Optional] [PJRT\_Plugin\_Initialize](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h#L173-L180). This is a one-time plugin setup, which will be called by the framework before any other functions are called.
+  * [Optional] [PJRT\_Plugin\_Attributes](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api.h#L182-L194).
+
+With the [wrapper](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api_wrapper_impl.h), you do not need to implement the remaining C APIs.
+
+
+### Step 2: Implement GetPjRtApi
+
+You need to implement a method `GetPjRtApi` which returns a `PJRT_Api*` containing function pointers to PJRT C API implementations. Below is an example assuming implementing through wrapper (similar to [pjrt\_c\_api\_cpu.cc](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api_cpu.cc)):
+
+```
+const PJRT_Api* GetPjrtApi() {
+  static const PJRT_Api pjrt_api =
+      pjrt::CreatePjrtApi(my_plugin::PJRT_Client_Create);
+  return &pjrt_api;
+}
+```
+
+### Step 3: Test C API implementations
+
+You can call [RegisterPjRtCApiTestFactory](https://github.com/openxla/xla/blob/c23fbd601a017be25726fd6d624b22daa6a8a4e5/xla/pjrt/c/pjrt_c_api_test.h#L31C6-L31C33) to run a small set of tests for basic PJRT C API behaviors.
+
+## How to use a PJRT plugin from JAX
+
+### Step 1: Set up JAX
+
+You can either use JAX nightly
+
+```
+pip install --pre -U jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
+
+pip install git+https://github.com/google/jax
+```
+or [build JAX from source](https://jax.readthedocs.io/en/latest/developer.html#building-jaxlib-from-source).
+
+For now, you need to match the jaxlib version with the PJRT C API version. It's usually sufficient to use a jaxlib nightly version from the same day as the TF commit you're building your plugin against, e.g.
+
+```
+pip install --pre -U jaxlib==0.6.1.dev20250428 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
+```
+You can also build a jaxlib from source at exactly the XLA commit you're building against ([instructions](https://jax.readthedocs.io/en/latest/developer.html#building-jaxlib-from-source-with-a-modified-xla-repository)).
+
+We will start supporting ABI compatibility soon.
+
+### Step 2: Use jax\_plugins namespace or set up entry\_point
+
+There are two options for your plugin to be discovered by JAX.
+
+1. Using namespace packages ([ref](https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/#using-naming-convention)). Define a globally unique module under the `jax_plugins` namespace package (i.e. just create a `jax_plugins` directory and define your module below it). Here is an example directory structure:
+
+```
+jax_plugins/
+  my_plugin/
+    __init__.py
+    my_plugin.so
+```
+2. Using package metadata ([ref](https://packaging.python.org/en/latest/guides/creating-and-discovering-plugins/#using-package-metadata)). If building a package via pyproject.toml or setup.py, advertise your plugin module name by including an entry-point under the `jax_plugins` group which points to your full module name. Here is an example via pyproject.toml or setup.py:
+
+```
+# use pyproject.toml
+[project.entry-points.'jax_plugins']
+my_plugin = 'my_plugin'
+
+# use setup.py
+entry_points={
+  "jax_plugins": [
+    "my_plugin = my_plugin",
+  ],
+}
+```
+Here are examples of how openxla-pjrt-plugin is implemented using Option 2: https://github.com/openxla/openxla-pjrt-plugin/pull/119, https://github.com/openxla/openxla-pjrt-plugin/pull/120.
+
+### Step 3: Implement an initialize() method
+
+You need to implement an initialize() method in your python module to register the plugin, for example:
+
+```
+import os
+import jax._src.xla_bridge as xb
+
+def initialize():
+  path = os.path.join(os.path.dirname(__file__), 'my_plugin.so')
+  xb.register_plugin('my_plugin', priority=500, library_path=path, options=None)
+```
+Please refer to [here](https://github.com/google/jax/blob/8f283bc9ed50d3828bd468ae57b1ee4df1527624/jax/_src/xla_bridge.py#L420) about how to use `xla_bridge.register_plugin`. It is currently a private method. A public API will be released in the future.
+
+You can run the line below to verify that the plugin is registered and raise an error if it can't be loaded.
+
+```
+jax.config.update("jax_platforms", "my_plugin")
+```
+JAX may have multiple backends/plugins. There are a few options to ensure your plugin is used as the default backend:
+*   Option 1: run `jax.config.update("jax_platforms", "my_plugin")` in the beginning of the program.
+*   Option 2: set ENV `JAX_PLATFORMS=my_plugin`.
+*   Option 3: set a high enough priority when calling xb.register\_plugin (the default value is 400 which is higher than other existing backends). Note the backend with highest priority will be used only when `JAX_PLATFORMS=''`. The default value of `JAX_PLATFORMS` is `''` but sometimes it will get overwritten.
+
+## How to test with JAX
+
+Some basic test cases to try:
+
+```
+# JAX 1+1
+print(jax.numpy.add(1, 1))
+# => 2
+
+# jit
+print(jax.jit(lambda x: x * 2)(1.))
+# => 2.0
+
+# pmap
+
+arr = jax.numpy.arange(jax.device_count()) print(jax.pmap(lambda x: x +
+jax.lax.psum(x, 'i'), axis_name='i')(arr))
+
+# single device: [0]
+
+# 4 devices: [6 7 8 9]
+```
+
+(We'll add instructions for running the jax unit tests against your plugin soon!)
+
+For more examples of PJRT plugins see [PJRT Examples](examples.md).