Update PJRT plugin API version to 0.57 (#19241)

It closes #19223.

`integrations/pjrt/third_party/pjrt_c_api/xla/pjrt/c/pjrt_c_api.h` is
updated to the latest (with API version from 0.38 to 0.57), fetching
from
https://github.com/openxla/xla/blob/a454e14ab0b10e35fb8ad73bd6db7d93782114f6/xla/pjrt/c/pjrt_c_api.h.

A blank implementation of `PJRT_Plugin_Attributes` is now provided since
an unimplemented `PJRT_Plugin_Attributes` will lead to initialization
failure of PJRT plugin (and thus crashes) in recent versions of PJRT
clients.

Also the JAX version in the CI workflow is updated from 0.4.20 to 0.4.35
and subsequently more tests can be enabled.

ci-exactly: build_packages, test_pjrt

---------

Signed-off-by: PragmaTwice <[email protected]>

Author: PragmaTwice

.github/workflows/pkgci_test_pjrt.yml

@@ -60,9 +60,8 @@ jobs:
           # install editable into venv
           source ${VENV_DIR}/bin/activate
           python -m pip install -v --no-deps -e integrations/pjrt/python_packages/iree_${{ matrix.pjrt_platform }}_plugin
-          # install jax (must be no larger than 0.4.20, refer to #19223)
-          # TODO: switch to the latest JAX after #19223 is fixed
-          python -m pip install jax==0.4.20 jaxlib==0.4.20 'numpy<2'
+          # install
+          python -m pip install jax==0.4.35
       - name: Run tests
         run: |
           source ${VENV_DIR}/bin/activate

build_tools/testing/run_jax_tests.sh

@@ -44,17 +44,9 @@ diff_jax_test() {
     echo "no difference found"
 }
 
-# FIXME: due to #19223, we need to use jax no higher than 0.4.20,
-# but in such version of jax, 'stablehlo.broadcast_in_dim' op
-# will be emitted without attribute 'broadcast_dimensions',
-# which leads to an error in IREE PJRT plugin.
-# So currently any program with broadcast will fail,
-# e.g. test/test_simple.py.
-# After #19223 is fixed, we can uncomment the line below.
-
-# diff_jax_test test/test_simple.py
-
 diff_jax_test test/test_add.py
+diff_jax_test test/test_degenerate.py
+diff_jax_test test/test_simple.py
 
 
 # FIXME: we can also utilize the native test cases from JAX,

integrations/pjrt/src/iree_pjrt/common/api_impl.cc

@@ -2164,6 +2164,15 @@ void BindMonomorphicApi(PJRT_Api* api) {
   BindUndefineds(api);
   ErrorInstance::BindApi(api);
 
+  // PJRT_Plugin_Attributes should be implemented since it will always be
+  // called from the PJRT client in the initial phase.
+  // here we provide a blank implementation to avoid crash due to unimplemented.
+  api->PJRT_Plugin_Attributes =
+      +[](PJRT_Plugin_Attributes_Args* args) -> PJRT_Error* {
+    args->num_attributes = 0;
+    args->attributes = nullptr;
+    return nullptr;
+  };
   api->PJRT_Plugin_Initialize =
       +[](PJRT_Plugin_Initialize_Args* args) -> PJRT_Error* { return nullptr; };
 

integrations/pjrt/src/iree_pjrt/common/stubs.inc

@@ -94,3 +94,8 @@
   _STUB(PJRT_Client_CreateViewOfDeviceBuffer);
   _STUB(PJRT_Executable_Fingerprint);
   _STUB(PJRT_Client_TopologyDescription);
+  _STUB(PJRT_Executable_GetCompiledMemoryStats);
+  _STUB(PJRT_Memory_Kind_Id);
+  _STUB(PJRT_ExecuteContext_Create);
+  _STUB(PJRT_ExecuteContext_Destroy);
+  _STUB(PJRT_Buffer_CopyRawToHost);

integrations/pjrt/third_party/pjrt_c_api/README.md

@@ -8,4 +8,4 @@ releases.
 Last synced from:
 
 * https://github.com/openxla/xla.git
-* commit: 96d1250d70c0bd6adf2778f31a266c1813fd107a
+* commit: a454e14ab0b10e35fb8ad73bd6db7d93782114f6