#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/pjrt_future.h"
#include "xla/pjrt/pjrt_stream_executor_client.h"
#include "xla/pjrt/tracked_device_buffer.h"
#include "xla/shape_util.h"
#include "xla/status_macros.h"
#include "xla/stream_executor/platform_manager.h"

    VLOG(2) << "Compiling using PJRT";
    Status status = CompileToPjRtLoadedExecutable(
        *ctx, platform_info_, function_, xla_compiler_args,
        DeviceCompileMode::kStrict, has_ref_vars_,
        /*may_alias_resource_update=*/true, &compilation_result, &pjrt_client,
        &pjrt_executable);
    OP_REQUIRES_OK_ASYNC(ctx, status, done);

    VLOG(2) << "Compiled using PJRT: " << status;

        "@llvm-project//mlir:IR",
    ],
)

cc_library(
    name = "tpu_passes",
    hdrs = ["transforms/tpu_passes.h"],
    visibility = [":friends"] + if_google([
        "//learning/brain/tfrt/ifrt/pjrt/__subpackages__",
        "//learning/serving/servables/tfrt:__subpackages__",
    ]),
    deps = [
        ":fallback_converter",
        ":tfrt_compile_options",
        "@llvm-project//llvm:Support",