#include "tensorflow/compiler/jit/xla_platform_info.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"

namespace tensorflow {

// Returns created options for the XLA compiler.
XlaCompiler::Options GenerateCompilerOptions(
    const DeviceCompiler<xla::LocalExecutable, xla::LocalClient>&
        xla_device_compiler,

        "@local_tsl//tsl/platform:status",
        "@local_xla//xla/stream_executor",
        "@local_xla//xla/stream_executor:executor_cache",
        "@local_xla//xla/stream_executor:host_memory_allocation",
        "@local_xla//xla/stream_executor:memory_allocation",
        "@local_xla//xla/stream_executor:platform",
        "@local_xla//xla/stream_executor:stream_executor_common",
    ],
)

cc_library(

#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/platform/mutex.h"

namespace tensorflow {
namespace device_compilation_cache_internal {
template <typename ExecutableType>
int64_t ExecutableSize(const ExecutableType* executable) {
  return 0;
}

template <>
inline int64_t ExecutableSize<xla::LocalExecutable>(

#include <string>

#include <gtest/gtest.h>
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "xla/xla_data.pb.h"

inline constexpr llvm::StringRef kXlaShardingAttrName = "_XlaSharding";

namespace tensorflow {
namespace {

TEST(DecodeShardingAttributeTest, CheckInvalidString) {
  xla::OpSharding sharding;
  EXPECT_TRUE(DecodeShardingAttribute("", sharding).succeeded());

#include "tensorflow/core/platform/status.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {

using PjRtDeviceCompiler =
    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>;

Status CompileToPjRtLoadedExecutable(
    const DeviceBase* device, const XlaPlatformInfo& platform_info,
    const NameAttrList& function,
    const std::vector<XlaCompiler::Argument>& args,

#ifndef TENSORFLOW_COMPILER_JIT_XLA_HOST_SEND_DEVICE_CONTEXT_H_
#define TENSORFLOW_COMPILER_JIT_XLA_HOST_SEND_DEVICE_CONTEXT_H_

#include "xla/shape.h"
#include "xla/stream_executor/device_memory.h"
#include "xla/stream_executor/stream.h"
#include "tensorflow/core/framework/device_base.h"
#include "tfrt/concurrency/async_value_ref.h"  // from @tf_runtime

namespace tensorflow {

#include "tensorflow/compiler/jit/pjrt_device_compiler_client.h"

#include <memory>
#include <string>
#include <utility>

namespace tensorflow {

xla::CompileOptions GetPjRtCompileOptions(
    const XlaCompiler::Options& options,
    const XlaCompiler::CompilationResult& result) {
  xla::CompileOptions pjrt_compile_options;
  pjrt_compile_options.argument_layouts = result.xla_input_shapes;
  pjrt_compile_options.executable_build_options =

    "@local_xla//xla:executable_run_options",
    "@local_xla//xla:status_macros",
    "@local_xla//xla:statusor",
    "@local_xla//xla/client:client_library",
    "@local_xla//xla/client:local_client",
    "@local_xla//xla/service:compiler",
    "@local_xla//xla/service/gpu:gpu_executable_run_options",
    "//tensorflow/core:core_cpu_internal",
    "//tensorflow/core:framework",
    "//tensorflow/core:lib",

    int32 count = 2;
  }

  // Describes a single XLA cluster.
  //
  // Next ID: 4
  message Cluster {
    string name = 1;

    // The number of nodes in the cluster.
    int32 size = 2;

    // A histogram of the TF operations in this cluster.
    repeated OpAndCount op_histogram = 3;
  }

  // The number of nodes in the graph that are not inside an XLA cluster.
  int32 unclustered_node_count = 1;

// xla_cluster_name: XLA cluster name for this XLA computation. We need it
//   because XLA cluster name might be different from `func_name`.
// func_name_attrs: they will be used to instantiate the XLA computation func.
// new_func_name: new function name for rewritten XLA computation func.
// host_compute_core: mapping from outside compilation cluster name to XLA
//   device assignment.
// fld: FunctionLibraryDefinition object.