#include "mlir/IR/Builders.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/protobuf/struct.pb.h"

namespace tensorflow {

using tsl::StatusOr;

// Converts an TensorFlow tensor proto into an MLIR elements attribute.

#include "tensorflow/c/tf_tensor.h"
#include "tensorflow/c/tf_tensor_helper.h"  // IWYU pragma: export
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/platform/casts.h"
#include "tensorflow/core/platform/status.h"

// Internal structures used by the C API. These are likely to change and should
// not be depended on.

#include "absl/status/status.h"
#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 {

#define TENSORFLOW_CC_SAVED_MODEL_LOADER_H_

#include <string>
#include <unordered_set>

#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/public/session.h"

namespace tensorflow {

/// Represents a SavedModel that is loaded from storage.
class SavedModelBundleInterface {

#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_UTILS_TEST_METADATA_CONFIG_H_

#include <variant>
#include <vector>

#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

// Fills in arg_shapes and metadata_proto with appropriate values based on the

#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h"
#include "tensorflow/core/lib/core/status.h"

namespace tensorflow {
namespace tf2xla {
namespace v2 {

// Run all the passes involved in transforming the graph before execution so

    mlir::tf_device::ParallelExecuteOp new_parallel_execute,
    mlir::OpBuilder* builder);

// Determines each logical core argument to metadata argument index mapping,
// based on sharding. The return value is indexed first by logical core then by
// argument index.
llvm::SmallVector<llvm::SmallVector<int64_t, 4>, 4> GetMetadataArgumentMapping(
    const tpu::TPUCompileMetadataProto& metadata);

#ifndef TENSORFLOW_CC_SAVED_MODEL_TEST_UTILS_H_
#define TENSORFLOW_CC_SAVED_MODEL_TEST_UTILS_H_

#include <ostream>
#include <string>

#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h"

namespace tensorflow::saved_model {

// TODO(b/229726259) Switch to OSS version after it's available.
// Simple implementation of a proto matcher comparing string representations.

#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_CLUSTER_TF_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_CLUSTER_TF_H_

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "tensorflow/core/lib/core/status.h"

namespace tensorflow {
namespace tf2xla {
namespace v1 {

// Run all the passes involved in transforming the graph before execution so

#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"

namespace mlir {
namespace TFL {

bool IsQuantized(const tflite::TensorT& tensor);

absl::StatusOr<mlir::quant::QuantizedType> GetQuantizedType(