test

pytorch · Dec 6, 2024 · 4fc0e0f · 4fc0e0f
1 parent 4c99d21
commit 4fc0e0f
Show file tree

Hide file tree

Showing 6 changed files with 15 additions and 14 deletions.
diff --git a/WORKSPACE b/WORKSPACE
@@ -50,7 +50,7 @@ new_local_repository(
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update the sha256 with the result.
 
-xla_hash = '30422ba2564eb108e12fb8a01629b2d0cf1f4393'
+xla_hash = '0020d6a56274c8893517adb22b46a3fcb5fa340a'
 
 http_archive(
     name = "xla",

diff --git a/setup.py b/setup.py
@@ -64,10 +64,11 @@
 
 base_dir = os.path.dirname(os.path.abspath(__file__))
 
-_date = '20241122'
-_libtpu_version = f'0.0.5.dev{_date}'
-_libtpu_storage_path = f'https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu/libtpu-{_libtpu_version}+nightly-py3-none-linux_x86_64.whl'
-_jax_version = f'0.4.36.dev{_date}'
+# libtpu 0.0.5 and JAX 0.4.36 are cut on Dec 5, 2024. If we're switching to a nightly build,
+# we should not regress to an earlier version.
+_libtpu_version = f'0.0.5'
+_libtpu_storage_path = f'https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu/libtpu-{_libtpu_version}-py3-none-linux_x86_64.whl'
+_jax_version = f'0.4.36'
 
 
 def _get_build_mode():

diff --git a/torch_xla/csrc/runtime/BUILD b/torch_xla/csrc/runtime/BUILD
@@ -162,7 +162,7 @@ cc_library(
     deps = [
         ":tf_logging",
         "@tsl//tsl/platform:stacktrace",
-        "@xla//xla:statusor",
+        "@tsl//tsl/platform:statusor",
     ],
 )
 
@@ -296,10 +296,10 @@ cc_library(
     deps = [
         ":tf_logging",
         ":profiler_backends",
+        "@com_google_absl//absl/status",
         "@xla//xla/backends/profiler/plugin:profiler_c_api_hdrs",
         "@xla//xla/backends/profiler/plugin:plugin_tracer",
         "@xla//xla/pjrt/c:pjrt_c_api_profiler_extension_hdrs",
-        "@xla//xla:status",
         "@tsl//tsl/profiler/lib:profiler_factory",
         "@xla//xla/tsl/profiler/rpc:profiler_server_impl",
         "@xla//xla/tsl/profiler/rpc/client:capture_profile",
@@ -380,7 +380,7 @@ cc_library(
     srcs = ["tf_logging.cc"],
     hdrs = ["tf_logging.h"],
     deps = [
-        "@xla//xla:statusor",
+        "@tsl//tsl/platform:statusor",
         "@xla//xla/service:platform_util",
         "@com_google_absl//absl/base:log_severity",
     ],
@@ -416,7 +416,7 @@ cc_library(
         "@com_google_absl//absl/types:optional",
         "@com_google_absl//absl/types:span",
         "@tsl//tsl/platform:hash",
-        "@xla//xla:statusor",
+        "@tsl//tsl/platform:statusor",
         "@xla//xla:types",
     ],
 )
@@ -479,17 +479,17 @@ ptxla_cc_test(
         ":computation_client",
         ":pjrt_computation_client",
         ":tensor_source",
+        "@com_google_absl//absl/status",
         "@xla//xla/tsl/lib/core:status_test_util",
         "@tsl//tsl/platform:env",
         "@tsl//tsl/platform:errors",
         "@tsl//tsl/platform:logging",
         "@tsl//tsl/platform:test",
         "@tsl//tsl/platform:test_main",
+        "@tsl//tsl/platform:statusor",
         "@xla//xla:literal",
         "@xla//xla:literal_util",
         "@xla//xla:shape_util",
-        "@xla//xla:status",
-        "@xla//xla:statusor",
         "@xla//xla/client:xla_builder",
         "@xla//xla/client:xla_computation",
         "@xla//xla/tests:literal_test_util",

diff --git a/torch_xla/csrc/runtime/debug_macros.h b/torch_xla/csrc/runtime/debug_macros.h
@@ -4,7 +4,7 @@
 #include "absl/status/status.h"
 #include "torch_xla/csrc/runtime/tf_logging.h"
 #include "tsl/platform/stacktrace.h"
-#include "xla/statusor.h"
+#include "tsl/platform/statusor.h"
 
 #define XLA_ERROR() TF_ERROR_STREAM()
 #define XLA_CHECK(c) TF_CHECK(c) << "\n" << tsl::CurrentStackTrace()

diff --git a/torch_xla/csrc/runtime/ifrt_computation_client_test.cc b/torch_xla/csrc/runtime/ifrt_computation_client_test.cc
@@ -11,12 +11,12 @@
 #include "torch_xla/csrc/runtime/tensor_source.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/logging.h"
+#include "tsl/platform/statusor.h"
 #include "tsl/platform/test.h"
 #include "xla/client/xla_builder.h"
 #include "xla/client/xla_computation.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
-#include "xla/statusor.h"
 #include "xla/tests/literal_test_util.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 

diff --git a/torch_xla/csrc/runtime/pjrt_computation_client_test.cc b/torch_xla/csrc/runtime/pjrt_computation_client_test.cc
@@ -12,12 +12,12 @@
 #include "torch_xla/csrc/runtime/tensor_source.h"
 #include "tsl/platform/env.h"
 #include "tsl/platform/logging.h"
+#include "tsl/platform/statusor.h"
 #include "tsl/platform/test.h"
 #include "xla/client/xla_builder.h"
 #include "xla/client/xla_computation.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
-#include "xla/statusor.h"
 #include "xla/tests/literal_test_util.h"
 #include "xla/tsl/lib/core/status_test_util.h"