pin update

[lsy323]

Accommodate the following changes:

• xla::Shape::rank() is renamed to xla::Shape::dimensions_size
• change of xla::Shape ctor
• PJRT compilation API change

[ysiraichi]

I was able to compile by adding the following patch to OpenXLA:

iff --git a/xla/service/gpu/model/gpu_collective_performance_model.cc b/xla/service/gpu/model/gpu_collective_performance_model.cc
index 496969f545..2d9f73ee36 100644
--- a/xla/service/gpu/model/gpu_collective_performance_model.cc
+++ b/xla/service/gpu/model/gpu_collective_performance_model.cc
@@ -34,7 +34,7 @@ limitations under the License.

 #if GOOGLE_CUDA
 #include "third_party/gpus/cuda/include/cuda.h"
-#include "third_party/gpus/cuda/nvml/include/nvml.h"
+#include "third_party/gpus/cuda/include/nvml.h"
 #endif  // GOOGLE_CUDA
 namespace xla {
 namespace gpu {
diff --git a/xla/service/gpu/model/gpu_collective_performance_model.h b/xla/service/gpu/model/gpu_collective_performance_model.h
index 01c3f3eb45..f44057602b 100644
--- a/xla/service/gpu/model/gpu_collective_performance_model.h
+++ b/xla/service/gpu/model/gpu_collective_performance_model.h
@@ -32,7 +32,7 @@ limitations under the License.
 #include <dlfcn.h>
 #endif

-#include "third_party/gpus/cuda/nvml/include/nvml.h"
+#include "third_party/gpus/cuda/include/nvml.h"
 // Below is a list of function pointers to be used
 // for querying device properties through nvml library.
 #define NVML_FUNCTOR(name, rettype, args) \

[lsy323]

I was able to compile by adding the following patch to OpenXLA:

iff --git a/xla/service/gpu/model/gpu_collective_performance_model.cc b/xla/service/gpu/model/gpu_collective_performance_model.cc
index 496969f545..2d9f73ee36 100644
--- a/xla/service/gpu/model/gpu_collective_performance_model.cc
+++ b/xla/service/gpu/model/gpu_collective_performance_model.cc
@@ -34,7 +34,7 @@ limitations under the License.

 #if GOOGLE_CUDA
 #include "third_party/gpus/cuda/include/cuda.h"
-#include "third_party/gpus/cuda/nvml/include/nvml.h"
+#include "third_party/gpus/cuda/include/nvml.h"
 #endif  // GOOGLE_CUDA
 namespace xla {
 namespace gpu {
diff --git a/xla/service/gpu/model/gpu_collective_performance_model.h b/xla/service/gpu/model/gpu_collective_performance_model.h
index 01c3f3eb45..f44057602b 100644
--- a/xla/service/gpu/model/gpu_collective_performance_model.h
+++ b/xla/service/gpu/model/gpu_collective_performance_model.h
@@ -32,7 +32,7 @@ limitations under the License.
 #include <dlfcn.h>
 #endif

-#include "third_party/gpus/cuda/nvml/include/nvml.h"
+#include "third_party/gpus/cuda/include/nvml.h"
 // Below is a list of function pointers to be used
 // for querying device properties through nvml library.
 #define NVML_FUNCTOR(name, rettype, args) \

Thank you @ysiraichi! I added this patch for now.

[lsy323]

Persistent cache test is failing on GPU, due to deserialization issue. Skipping the test for now and will file a Github Issue for this.

[ RUN      ] PersistentCacheTest.test_persistent_cache_mp
E0402 19:38:44.492582203   21190 server_chttp2.cc:40]        ***"created":"@1743622724.492561846","description":"Only 1 addresses added out of total 2 resolved","file":"external/com_github_grpc_grpc/src/core/ext/transport/chttp2/server/chttp2_server.cc","file_line":404,"referenced_errors":[***"created":"@1743622724.492558996","description":"Address family not supported by protocol","errno":97,"file":"external/com_github_grpc_grpc/src/core/lib/iomgr/socket_utils_common_posix.cc","file_line":420,"os_error":"Address family not supported by protocol","syscall":"socket","target_address":"[::1]:8547"***]***
E0402 19:38:58.918947906   22327 server_chttp2.cc:40]        ***"created":"@1743622738.918928838","description":"Only 1 addresses added out of total 2 resolved","file":"external/com_github_grpc_grpc/src/core/ext/transport/chttp2/server/chttp2_server.cc","file_line":404,"referenced_errors":[***"created":"@1743622738.918926006","description":"Address family not supported by protocol","errno":97,"file":"external/com_github_grpc_grpc/src/core/lib/iomgr/socket_utils_common_posix.cc","file_line":420,"os_error":"Address family not supported by protocol","syscall":"socket","target_address":"[::1]:8547"***]***
2025-04-02 19:39:05.879197: W torch_xla/csrc/runtime/pjrt_computation_client.cc:680] Failed to deserialize executable: UNIMPLEMENTED: Deserializing serialized executable not supported.
2025-04-02 19:39:05.880553: W torch_xla/csrc/runtime/pjrt_computation_client.cc:680] Failed to deserialize executable: UNIMPLEMENTED: Deserializing serialized executable not supported.
2025-04-02 19:39:05.887184: W torch_xla/csrc/runtime/pjrt_computation_client.cc:680] Failed to deserialize executable: UNIMPLEMENTED: Deserializing serialized executable not supported.
2025-04-02 19:39:05.892825: W torch_xla/csrc/runtime/pjrt_computation_client.cc:680] Failed to deserialize executable: UNIMPLEMENTED: Deserializing serialized executable not supported.

[ManfeiBai]

Thanks for the amazing work, its a really huge change adopt PR, LGTM

[vanbasten23]

Thanks @lsy323 for updating the pin.

Regarding the paged_attention hang, could you update this line

xla/torch_xla/experimental/custom_kernel.py

Line 1212 in c044c69

step = torch.zeros((1,), dtype=torch.int32).to("xla")

to step = torch.ones((1,), dtype=torch.int32).to("xla")? It should make the test pass. I tested locally.

[lsy323]

Thanks @lsy323 for updating the pin.

Regarding the paged_attention hang, could you update this line

xla/torch_xla/experimental/custom_kernel.py

Line 1212 in c044c69

step = torch.zeros((1,), dtype=torch.int32).to("xla")

to step = torch.ones((1,), dtype=torch.int32).to("xla")? It should make the test pass. I tested locally.

Thanks @vanbasten23! Updated the PR. Also do you mind elaborating a bit on this?

[vanbasten23]

Thanks @lsy323 for updating the pin.
Regarding the paged_attention hang, could you update this line

xla/torch_xla/experimental/custom_kernel.py

Line 1212 in c044c69

step = torch.zeros((1,), dtype=torch.int32).to("xla")

to step = torch.ones((1,), dtype=torch.int32).to("xla")? It should make the test pass. I tested locally.

Thanks @vanbasten23! Updated the PR. Also do you mind elaborating a bit on this?

Yeah, jax-ml/jax@8c73799 made a change (it's not a bug but a valid change). As a result, the torch_xla wrapper needs to change accordingly.