#ifndef TENSORFLOW_COMPILER_MLIR_TFR_INTEGRATION_NODE_EXPANSION_PASS_H_
#define TENSORFLOW_COMPILER_MLIR_TFR_INTEGRATION_NODE_EXPANSION_PASS_H_

#include "tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.h"
#include "tensorflow/core/common_runtime/eager/eager_op_rewrite_registry.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace tfr {

// An optimization pass that decompose the composite ops in a module according

#include "mlir/IR/Types.h"  // from @llvm-project
#include "mlir/IR/Value.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
#include "xla/statusor.h"
#include "tensorflow/core/framework/types.pb.h"

namespace mlir {
class Builder;
}  // namespace mlir

namespace tflite {
// Convert the MLIR type to the corresponding TFLite tensor.

#define TENSORFLOW_CC_SAVED_MODEL_IMAGE_FORMAT_INTERNAL_API_H_

#include <string>
#include <tuple>

#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"

#define IS_OSS false

namespace tensorflow {
namespace image_format {

// Reads the SavedModel proto from {file_prefix}{.pb|.cpb}.

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

#ifndef TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_H_
#define TENSORFLOW_COMPILER_JIT_DEADNESS_ANALYSIS_H_

#include "tensorflow/core/graph/graph.h"

namespace tensorflow {

// This analyzes a TensorFlow graph to identify nodes which may have partially
// dead inputs (i.e. these nodes may have some dead inputs and some alive
// inputs).
//

std::unique_ptr<OperationPass<ModuleOp>> CreateCastBf16OpsToF32Pass();

// Creates a pass that lifts HashTable ops as function arguments. In the graph
// execution mode, resource ops with the same `shared_name` attribute point to
// the same underlying resource. This is not true in the eager execution mode.
// Lifting resource ops as arguments will help unifying them across functions.
std::unique_ptr<OperationPass<ModuleOp>> CreateLiftHashTableOpsAsArgsPass();

#include "tensorflow/compiler/jit/pjrt_base_device.h"
#include "tensorflow/compiler/jit/xla_device.h"
#include "xla/stream_executor/integrations/tf_allocator_adapter.h"
#include "tensorflow/core/framework/op_kernel.h"

namespace tensorflow {

// Holds some information about the platform on which an
// XlaLaunch/_XlaCompile/_XlaRun op must run on. Provides a common layer of

//
// The custom device defines copy operations for moving TensorHandles on and
// off, and an execution operation for named operations. Often execution will
// simply wrap op execution on one or more physical devices.
//
// device_info is an opaque caller-defined type stored with the custom device
// which is passed to the functions referenced in the TFE_CustomDevice struct

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

#include "absl/container/flat_hash_set.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/lite/tools/optimize/reduced_precision_support.h"

namespace mlir {
namespace quant {

// Stores information about how to quantize a user-specified custom operation.

    llvm::ArrayRef<int64_t> device_assignment_attr);

// Virtual device name of the passed logical core. The logical core is the index
// of a core within a replica.
std::string GetDeviceAliasForLogicalCore(int core_index);

// Virtual device name of the host that is associated with the passed logical
// core. The logical core is the index of a core within a replica.
std::string GetDeviceAliasForHostOfLogicalCore(int core_index);

#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/protobuf/trackable_object_graph.pb.h"