#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"
#include "tensorflow/core/lib/core/errors.h"

constexpr absl::string_view kGroupKeyAttrName =
    "tf2xla.collective_info.group_key";

// Extracts shape from XlaArgument as TensorShape. If shape is a xla::Shape,
// that is converted to a TensorShape.
absl::StatusOr<TensorShape> GetTensorShapeFromXlaArgument(
    const XlaArgument& arg) {
  if (absl::holds_alternative<xla::Shape>(arg.shape)) {
    TensorShape arg_shape;
    TF_RETURN_IF_ERROR(

  Status AddInstructionNode(Operation* inst);
  void UseOriginalFunctionNames(NodeDef& node_def);

  Status AddEdge(Operation* inst);

  absl::StatusOr<std::unique_ptr<NodeDef>> GetArgumentNode(
      BlockArgument arg, unsigned index, llvm::StringRef name);
  absl::StatusOr<std::unique_ptr<NodeDef>> GetReturnNode(FuncOp function,
                                                         Value operand,

  Status AddInstructionNode(Operation* inst);
  void UseOriginalFunctionNames(NodeDef& node_def);

  Status AddEdge(Operation* inst);

  absl::StatusOr<std::unique_ptr<NodeDef>> GetArgumentNode(
      BlockArgument arg, unsigned index, llvm::StringRef name);
  absl::StatusOr<std::unique_ptr<NodeDef>> GetReturnNode(FuncOp function,
                                                         Value operand,

#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/lib/core/status_test_util.h"
#include "tsl/lib/monitoring/test_utils.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace tf2xla {
namespace v2 {

using ::tensorflow::monitoring::testing::CellReader;
using ::testing::Not;
using ::testing::TestWithParam;
using tpu::FunctionToHloArgs;

#include <optional>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>

#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "llvm/Support/ToolOutputFile.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/OwningOpRef.h"  // from @llvm-project

==============================================================================*/

#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>

#include "absl/log/log.h"
#include "absl/status/statusor.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"

namespace mlir {
namespace TFL {
namespace {
#define GEN_PASS_DEF_LIFTTFLITEFLEXOPSPASS
#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

using ::tensorflow::StatusOr;

constexpr StringRef kFlexOpNamePrefix = "Flex";

// Pattern that converts TFL::CustomOp that encodes a Flex op into a TF dialect
// operation.
class LiftFlexCustomOp : public OpRewritePattern<TFL::CustomOp> {

#include "tensorflow/compiler/mlir/tensorflow/utils/xla_call_module_attrs.h"
#include "tensorflow/compiler/tf2xla/kernels/xla_call_module_loader.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"  // IWYU pragma: keep
#include "tsl/platform/statusor.h"

namespace mlir {
namespace TF {
namespace {

#define GEN_PASS_DEF_XLACALLMODULEDESERIALIZATIONPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tensorflow/core/tpu/tpu_compile.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/lib/monitoring/sampler.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/status.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {