#include "xla/xla_data.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace tensorflow {

inline constexpr llvm::StringRef kInputShardingAttr =
    "input_sharding_configuration";
inline constexpr llvm::StringRef kOutputShardingAttr =
    "output_sharding_configuration";

// Parses the sharding string. This sharding string can be binary (serialized)
// or human readable.

namespace {

constexpr StringRef kStablehloModuleAttrsAttrName = "_stablehlo_module_attrs";
constexpr StringRef kUsesShapePolymorphismAttr = "jax.uses_shape_polymorphism";

// Default version number for native serialization.
constexpr int64_t kDefaultVersion = 9;
// Platforms for XlaCallModuleOp.
constexpr StringRef kPlatformCpu = "CPU";
constexpr StringRef kPlatformTpu = "TPU";

#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"

namespace mlir {
namespace TFL {
namespace tac {

constexpr char ExampleHardware::kId[];  // Define kId.

mlir::RewritePatternSet ExampleHardware::GetTransformations(
    MLIRContext* context) const {
  mlir::RewritePatternSet patterns(context);

#include "tensorflow/compiler/mlir/tensorflow/utils/xla_sharding_util.h"

#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;

// both implement "func_1".
//
// The attribute on the FuncOp means what it actually implements while the
// attribute on the CallOp means what it actually looks for.
constexpr char kInterfaceNameAttr[] = "tac.interface_name";

inline std::optional<std::string> GetInterFaceName(Operation* op) {
  auto name_attr = op->getAttrOfType<StringAttr>(kInterfaceNameAttr);
  if (!name_attr) return std::nullopt;

#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/core/lib/io/path.h"

namespace mlir {
namespace TF {
namespace {

static constexpr int kTextFileIndex_WholeLine = -2;
static constexpr int kTextFileIndex_LineNumber = -1;

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

using ::mlir::func::FuncOp;
using ::tsl::Status;
using ::tsl::StatusOr;
using ::xla::ReplicaGroup;
using ::xla::ShapeUtil;
using ::xla::XlaBuilder;
using ::xla::XlaComputation;
using ::xla::XlaOp;

static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1442 : i32}} {

    {tsl::monitoring::Buckets::Exponential(1, 1.5, 45)});

// Name of component for error logging. This name is fixed and required to
// enable logging.
constexpr char kBridgeComponent[] = "TFXLABridge";
constexpr char kFullBridge[] = "full_bridge";

namespace {

bool ShouldFallbackToGraphCompiler(
    const std::variant<MlirToHloArgs, FunctionToHloArgs>& computation) {

#include "tensorflow/core/util/device_name_utils.h"

namespace tensorflow {
using tsl::StatusOr;

inline constexpr absl::string_view kNumCoresPerReplicaAttr =
    "num_cores_per_replica";
inline constexpr absl::string_view kTopologyAttr = "topology";
inline constexpr absl::string_view kDeviceAssignmentAttr = "device_assignment";

// A TPU device for execution alongside its associated host CPU device.

#include <string>
#include <utility>

#include "llvm/ADT/StringExtras.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project

namespace mlir {
namespace quant {

constexpr StringRef kTfEntryFunctionAttr = "tf.entry_function";

void AddEntryFunctionInput(StringRef input_name, func::FuncOp func_op) {
  auto entry_func_attr =
      func_op->getAttrOfType<DictionaryAttr>(kTfEntryFunctionAttr);