#include "tensorflow/compiler/jit/xla_compilation_cache.pb.h"
#include "tensorflow/compiler/jit/xla_device_compiler_client.h"
#include "xla/client/client_library.h"
#include "xla/client/executable_build_options.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/tfrt_cpu_pjrt_client.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"

namespace {
using ::testing::_;
using ::testing::Return;

using XlaDeviceCompiler =
    DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
using XlaDeviceExecutablePersistor =
    DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient>;
using Signature = DeviceCompilationClusterSignature;

xla::LocalClient* GetLocalClient() {
  // TODO(b/255826209): Figure out how to run this test with the CPU client as
  // well.

#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/layout.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/shape.h"
#include "xla/stream_executor/tpu/c_api_conversions.h"
#include "xla/stream_executor/tpu/tpu_api.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"

namespace mlir {

static FailureOr<std::vector<int64_t>> GetTPUInfeedLayoutFromAPI(

        "@local_xla//xla:shape_util",
        "@local_xla//xla/pjrt:pjrt_api",
        "@local_xla//xla/pjrt:pjrt_c_api_client",
        "@local_xla//xla/pjrt/c:pjrt_c_api_cpu",
        "@local_xla//xla/pjrt/c:pjrt_c_api_hdrs",
        "@local_xla//xla/pjrt/c:pjrt_c_api_wrapper_impl",
        "@local_xla//xla/pjrt/cpu:cpu_client",
    ],

#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/client/xla_builder.h"
#include "xla/client/xla_computation.h"
#include "xla/hlo/ir/hlo_instruction.h"
#include "xla/hlo/ir/hlo_opcode.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/service/hlo.pb.h"
#include "xla/translate/hlo_to_mhlo/hlo_function_importer.h"
#include "xla/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"

#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
#include "xla/client/client_library.h"
#include "xla/client/compile_only_client.h"
#include "xla/client/xla_computation.h"
#include "xla/service/cpu/cpu_compiler.h"
#include "xla/statusor.h"
#include "xla/stream_executor/platform_manager.h"
#include "xla/util.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/graph.pb.h"

  }

  return mlir::success();
}

bool UnsupportedPartitionedShardingType(xla::OpSharding::Type sharding) {
  return sharding != xla::OpSharding::REPLICATED &&
         sharding != xla::OpSharding::OTHER;
}

}  // namespace

absl::StatusOr<std::map<int, int>> GetDimensionIndicesAndNumSplitsFromSharding(
    const xla::OpSharding& sharding) {
  int64_t tensor_tile_rank = sharding.tile_assignment_dimensions_size();

#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/core/status.h"

namespace tensorflow {

// An OpKernel that compiles an op to an XLA computation and runs it. Unlike

// Returns the XLA cluster in which `node` is placed if it is in an XLA cluster,
// otherwise returns nullopt.
std::optional<absl::string_view> GetXlaClusterForNode(const Node& node);

// Removes `node_def` its XLA cluster (by clearing its _XlaCluster attribute).
void RemoveFromXlaCluster(NodeDef* node_def);

// Removes `node` its XLA cluster (by clearing its _XlaCluster attribute).

#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/xla_sharding_util.h"
#include "xla/xla.pb.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status.h"