#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/xla_computation.h"
#include "xla/hlo/ir/hlo_sharding.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "xla/shape.h"
#include "xla/translate/mhlo_to_hlo/layout_util.h"
#include "xla/translate/mhlo_to_hlo/mlir_hlo_to_hlo.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"

#include <vector>

#include <gtest/gtest.h>
#include "absl/log/check.h"
#include "xla/pjrt/c/pjrt_c_api.h"
#include "xla/pjrt/c/pjrt_c_api_cpu.h"
#include "xla/pjrt/c/pjrt_c_api_wrapper_impl.h"
#include "xla/pjrt/cpu/cpu_client.h"
#include "xla/pjrt/pjrt_api.h"
#include "xla/pjrt/pjrt_c_api_client.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "tensorflow/core/framework/types.h"

mlir::LogicalResult DecodeShardingAttribute(const std::string& shard_str,
                                            xla::OpSharding& sharding,
                                            bool report_error = true);

// Encodes the sharding in human readable form.
mlir::LogicalResult DecodeShardingAttribute(mlir::Attribute shard_attr,
                                            xla::OpSharding& sharding,
                                            bool report_error = true);

        "@local_xla//xla/client:padding",
        "@local_xla//xla/client:sharding_builder",
        "@local_xla//xla/mlir_hlo",
        "@local_xla//xla/mlir_hlo:convert_op_folder",
        "@local_xla//xla/mlir_hlo:mhlo_passes",
        "@local_xla//xla/mlir_hlo:type_conversion",
        "@local_xla//xla/stream_executor/tpu:c_api_conversions",
        "@local_xla//xla/stream_executor/tpu:tpu_api",

        "@local_xla//xla:status_macros",
        "@local_xla//xla/client:local_client",
        "@local_xla//xla/pjrt:pjrt_client",
        "@local_xla//xla/pjrt:pjrt_future",
        "@local_xla//xla/pjrt:pjrt_stream_executor_client",
        "@local_xla//xla/pjrt:tracked_device_buffer",
        "@local_xla//xla/service:shaped_buffer",
        "@local_xla//xla/stream_executor:device_memory_allocator",

      op->emitWarning("TF/XLA TPU bridge input check: found invalid op. ")
          << op->getName() << " can't be both xla and non-xla";
      return true;
    }
    op->emitOpError("TF/XLA TPU bridge input check: found invalid op. ")
        << "Can't be both xla and non-xla";
    return false;
  }
  return true;
}
void GetFlattenedShardings(llvm::SmallVector<xla::OpSharding>& shardings_result,

#include "absl/types/variant.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"

namespace tensorflow {
namespace tf2xla {
namespace v1 {

// Compiles the given Tensorflow graph into xla::HLO. The result is in

#include "xla/pjrt/c/pjrt_c_api.h"
#include "xla/pjrt/pjrt_c_api_client.h"
#include "tensorflow/core/framework/tensor.h"

namespace tensorflow {

absl::StatusOr<PJRT_Buffer*> GetPjRtCBufferFromTensor(const Tensor* tensor);

absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer,
                                    xla::PjRtCApiClient* c_api_client,

    const XlaPlatformInfo& platform_info, DeviceType compilation_device_type,
    DeviceCompiler<xla::LocalExecutable, xla::LocalClient>**
        xla_device_compiler);

// Fetches a DeviceCompiler from the tfrt_global resource manager (or creates
// one there if not found) that uses xla::PjRtClient using an appropriate
// PjRtClient for `platform_info.device_type()` and sets *pjrt_device_compiler

    .Output("results: Tresults")
    .Attr("Tresults: list(type) >= 0")
    .Attr("function: func")
    // XLA random-number generation ops are stateful.
    // TODO(phawkins): create stateful and non-stateful variants of XlaLaunch.
    .SetIsStateful()
    .Doc("XLA Launch Op. For use by the XLA JIT only.");

REGISTER_OP("XlaLaunchV2")
    .Input("args: Targs")
    .Output("results: Tresults")