diff --git a/WORKSPACE b/WORKSPACE index 7106673a25b..3869f82a3d3 100644 --- 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 index e4db95ef14b..827c31f14c6 100644 --- 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 index 0783a2d98ac..daf93bda365 100644 --- 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 index 8eca8034e34..70d87321f7f 100644 --- 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 index 4027f902e6b..8f05c1cc819 100644 --- 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 index be42a5e4dbf..85f28edcc3e 100644 --- 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"