#include "mlir/IR/Operation.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
#include "tfrt/core_runtime/tensor_handle.h"  // from @tf_runtime

namespace tfrt {

#include <string>

#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/StringMap.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status.h"

namespace tensorflow {

struct ArrayInfoBase {
  // The node type when the input node is imported. Typically needs to be

#include <vector>

#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/jit/xla_activity_listener.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/test.h"

namespace tensorflow {

// A listener to inspect the use of XLA's persistent compilation cache entries.

#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/lite/toco/model_flags.pb.h"
#include "tensorflow/lite/toco/toco_flags.pb.h"
#include "tensorflow/lite/toco/types.pb.h"

#include <functional>
#include <iostream>
#include <ostream>

#include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
#include "xla/client/xla_computation.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"

namespace tensorflow {
namespace tfcompile {

using QuantizeXlaFn = std::function<Status(const tf2xla::Config& config,

#define TENSORFLOW_COMPILER_JIT_PJRT_DEVICE_CONTEXT_H_

#include <utility>

#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/platform/status.h"

namespace tensorflow {

// Helper class for managing data transfers between host and accelerator
// devices using PjRt.
class PjRtDeviceContext : public DeviceContext {

#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/platform/statusor.h"

namespace tensorflow {

// Populates the tf.versions attribute on a module, given a corresponding
// graph VersionDef proto.

#include <string>

#include "tensorflow/compiler/aot/flags.h"
#include "tensorflow/compiler/tf2xla/tf2xla.pb.h"
#include "xla/service/cpu/cpu_compiler.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/platform/status.h"

namespace tensorflow {
namespace tfcompile {

// CompileResult describes the output of CompileGraph, where the object file
// data and meta-information is available in aot.

#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saver.pb.h"

namespace mlir::quant::stablehlo {

// Suffix string for the module export step. Used for debugging.

#define TENSORFLOW_COMPILER_JIT_PJRT_BASE_DEVICE_H_

#include <string>
#include <utility>
#include <vector>

#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/core/common_runtime/local_device.h"
#include "tensorflow/core/framework/device_base.h"

namespace tensorflow {

// tensorflow::PjRtBaseDevice replaces the deprecated tensorflow::XlaDevice.
// This accelerator agnostic device is mainly used to store metadata.