Skip to content

Commit

Permalink
[Pin Update] Version 20230826(#5527) (#5555)
Browse files Browse the repository at this point in the history
Co-authored-by: Jiewen Tan <[email protected]>
  • Loading branch information
ManfeiBai and alanwaketan authored Sep 14, 2023
1 parent 9bcf3f7 commit 2c07df9
Show file tree
Hide file tree
Showing 4 changed files with 7 additions and 7 deletions.
4 changes: 2 additions & 2 deletions WORKSPACE
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,9 @@ http_archive(
"//openxla_patches:gpu_race_condition.diff",
"//openxla_patches:constexpr_return.diff",
],
strip_prefix = "xla-cd2cf5c34931e4fc1cacf83bfc480a5b93f05f6d",
strip_prefix = "xla-7a371ed44aba34f83d6d3d1159d2e6d0d327c603",
urls = [
"https://github.com/openxla/xla/archive/cd2cf5c34931e4fc1cacf83bfc480a5b93f05f6d.tar.gz",
"https://github.com/openxla/xla/archive/7a371ed44aba34f83d6d3d1159d2e6d0d327c603.tar.gz",
],
)

Expand Down
3 changes: 2 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,8 @@

base_dir = os.path.dirname(os.path.abspath(__file__))

_libtpu_version = '0.1.dev20230809'
_libtpu_version = '0.1.dev20230826'
_libtpu_storage_path = f'https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/wheels/libtpu-nightly/libtpu_nightly-{_libtpu_version}-py3-none-any.whl'


def _get_build_mode():
Expand Down
4 changes: 2 additions & 2 deletions test/cpp/run_tests.sh
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ fi


if [ "$LOGFILE" != "" ]; then
bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all ${FILTER:+"$FILTER"} 2> $LOGFILE
bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all --test_timeout 1000 ${FILTER:+"$FILTER"} 2> $LOGFILE
else
bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all ${FILTER:+"$FILTER"}
bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all --test_timeout 1000 ${FILTER:+"$FILTER"}
fi
3 changes: 1 addition & 2 deletions torch_xla/csrc/runtime/pjrt_computation_client.cc
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@
#include "xla/pjrt/tfrt_cpu_pjrt_client.h"
#include "xla/pjrt/tpu_client.h"
#include "xla/shape.h"
#include "xla/stream_executor/tpu/tpu_initializer_framework_helper.h"

using xla::internal::XlaBuilderFriend;

Expand Down Expand Up @@ -105,7 +104,7 @@ PjRtComputationClient::PjRtComputationClient() {
TF_VLOG(1) << "Initializing TFRT TPU client...";
XLA_CHECK_OK(pjrt::LoadPjrtPlugin(
"tpu", sys_util::GetEnvString(env::kEnvTpuLibraryPath, "libtpu.so")));
tsl::Status tpu_status = tensorflow::tpu::FindAndLoadTpuLibrary();
tsl::Status tpu_status = pjrt::InitializePjrtPlugin("tpu");
XLA_CHECK(tpu_status.ok());
client_ = std::move(xla::GetCApiClient("TPU").value());
} else if (device_type == "TPU_LEGACY") {
Expand Down

0 comments on commit 2c07df9

Please sign in to comment.