Update the tf pin to 2023 01/23 (#4474)

* Update the tf pin to 2023 01/18

* Update CAPI path

* update tf pin

* Update libtpu

* Update libtpu again to 01/24

Co-authored-by: root <root@t1v-n-7513beeb-w-0.us-central1-a.c.tpu-pytorch.internal>

Author: JackCaoG

setup.py

@@ -60,7 +60,7 @@
 base_dir = os.path.dirname(os.path.abspath(__file__))
 third_party_path = os.path.join(base_dir, 'third_party')
 
-_libtpu_version = '0.1.dev20221223'
+_libtpu_version = '0.1.dev20230124'
 _libtpu_storage_path = f'https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/wheels/libtpu-nightly/libtpu_nightly-{_libtpu_version}-py3-none-any.whl'
 
 

third_party/tensorflow

@@ -15,7 +15,7 @@
 #include "tensorflow/compiler/xla/pjrt/tfrt_cpu_pjrt_client.h"
 #include "tensorflow/compiler/xla/pjrt/tpu_client.h"
 #include "tensorflow/compiler/xla/shape.h"
-#include "tensorflow/compiler/xla/stream_executor/tpu/pjrt_api.h"
+#include "tensorflow/compiler/xla/pjrt/pjrt_api.h"
 #include "tensorflow/compiler/xla/xla_client/computation_client.h"
 #include "tensorflow/compiler/xla/xla_client/debug_macros.h"
 #include "tensorflow/compiler/xla/xla_client/env_vars.h"
@@ -61,7 +61,7 @@ PjRtComputationClient::PjRtComputationClient() {
     client_ = xla::GetTpuClient(max_inflight_computations).value();
   } else if (device_type == "TPU_C_API") {
     TF_VLOG(1) << "Initializing PjRt C API client...";
-    XLA_CHECK_OK(stream_executor::tpu::LoadPjrtPlugin(
+    XLA_CHECK_OK(pjrt::LoadPjrtPlugin(
         "tpu", sys_util::GetEnvString(env::kEnvTpuLibraryPath, "libtpu.so")));
     client_ = std::move(xla::GetCApiClient("TPU").value());
   } else if (device_type == "GPU") {

third_party/xla_client/pjrt_computation_client.cc

@@ -14,7 +14,6 @@
 #include "absl/strings/str_split.h"
 #include "tensorflow/cc/ops/const_op.h"
 #include "tensorflow/compiler/xla/shape_util.h"
-#include "tensorflow/compiler/xla/stream_executor/lib/mathutil.h"
 #include "tensorflow/compiler/xla/util.h"
 #include "tensorflow/compiler/xla/xla_client/env_vars.h"
 #include "tensorflow/compiler/xla/xla_client/multi_wait.h"
@@ -27,6 +26,7 @@
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/util/device_name_utils.h"
+#include "tensorflow/tsl/lib/math/math_util.h"
 
 namespace xla {
 namespace {
@@ -76,7 +76,7 @@ class TensorAllocator : public tensorflow::Allocator {
     alignment = std::max<size_t>(alignment, sizeof(void*));
     // To call aligned_alloc(), num_bytes must be multiple of alignment.
     num_bytes =
-        stream_executor::port::MathUtil::CeilOfRatio(num_bytes, alignment) * alignment;
+        tsl::MathUtil::CeilOfRatio(num_bytes, alignment) * alignment;
 
     AllocKey alloc_key = {alignment, num_bytes};
     void* block = nullptr;