#include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h"
#include "tensorflow/compiler/tf2xla/xla_context.h"
#include "tensorflow/compiler/tf2xla/xla_expression.h"
#include "xla/client/xla_builder.h"
#include "xla/client/xla_computation.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/op_kernel.h"

namespace mlir {
namespace mhlo {

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

#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/framework/function.h"
#include "tsl/framework/device_id_utils.h"

namespace tensorflow {
namespace {
using XlaDeviceCompiler =
    DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
using PjRtDeviceCompiler =
    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>;

#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/pjrt/compile_options.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/tpu/kernels/tpu_compile.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"

#include "llvm/MC/TargetRegistry.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Target/TargetOptions.h"
#include "llvm/TargetParser/Triple.h"
#include "xla/service/llvm_ir/llvm_type_conversion_util.h"
#include "xla/util.h"

namespace tensorflow {
namespace tfcompile {

using xla::llvm_ir::AsStringRef;

static void AddEmbeddedProtocolBufferToLlvmModule(
    llvm::Module* module, const ::tensorflow::protobuf::MessageLite& proto,

        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:ArithDialect",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Support",
        "@llvm-project//mlir:TransformUtils",
        "@local_xla//xla/mlir_hlo",
    ],
)

cc_library(
    name = "scatter",
    srcs = [
        "scatter.cc",
    ],
    hdrs = [
        "scatter.h",
    ],
    deps = [
        ":util",

#include <optional>

#include "mlir/IR/Diagnostics.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "xla/xla.pb.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace mlir {

  std::vector<int32_t> data{1, 2, 3, 4, 5, 6};
  xla::Shape xla_shape = xla::ShapeUtil::MakeShape(xla::S32, dimensions);
  TF_ASSERT_OK_AND_ASSIGN(
      auto pjrt_buffer,
      pjrt_client->BufferFromHostBuffer(
          data.data(), xla_shape.element_type(), xla_shape.dimensions(),
          /*byte_strides=*/std::nullopt,
          xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,

def LegalizeTFCommunicationPass : Pass<"xla-legalize-tf-communication", "ModuleOp"> {
  let summary = "Legalize TF/XLA communication ops (TensorFlow dialect) to the HLO "
           "dialect";
  let dependentDialects = ["mhlo::MhloDialect", "sparse_tensor::SparseTensorDialect"];
  let constructor = "::mlir::mhlo::CreateLegalizeTFCommunicationPass()";
  let description = [{
    A pass that legalizes TF/XLA communication ops, propagates their respective

#include "tensorflow/compiler/mlir/tensorflow/transforms/set_tpu_infeed_layout.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 {
namespace mhlo {
namespace {

// RUN: tf-mlir-translate -mlir-tf-to-hlo-text %s -tf-input-shapes=8,16,16,64:64 -tf-xla-emit-use-tuple-args -tf-xla-emit-return-tuple | FileCheck %s
// RUN: tf-mlir-translate -mlir-tf-to-hlo-text-via-builder %s -tf-input-shapes=8,16,16,64:64 | FileCheck %s

module attributes {tf.versions = {producer = 179 : i32}} {