#include "absl/strings/string_view.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"

namespace stablehlo::quantization::pywrap {

// Function used by the pywrap_quantization module to mirror
// `::mlir::quant::stablehlo::QuantizeStaticRangePtq`.

#include <string>
#include <vector>

#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

// Compiles a serialized MLIR module and returns a serialized MLIR module of the

#ifndef TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_DISTRIBUTED_MANAGER_H_
#define TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_DISTRIBUTED_MANAGER_H_

#include <cstdint>
#include <string>

#include "tensorflow/core/platform/status.h"

namespace tsl {
class CoordinationServiceAgent;
}

namespace tensorflow {
class ImmediateExecutionContext;
class ServerDef;
class WorkerEnv;
class WorkerCacheInterface;

DenseIntElementsAttr DenseI64AttrToI32Attr(
    const DenseIntElementsAttr& dense_attr, PatternRewriter& builder);

// Returns a NHWC shaped type from an NCHW shaped type op.
// For example- Given a Composite op that wraps a core.aten.avg_pool2d, this
// returns the return type of the tfl.average_pool_2d emitted. Note that the
// aten.avg_pool2d works with the NCHW layout while tfl.average_pool_2d assumes
// NHWC.

#define TENSORFLOW_COMPILER_JIT_XLA_HOST_RECV_DEVICE_CONTEXT_H_

#include "xla/shape.h"
#include "xla/stream_executor/device_memory.h"
#include "xla/stream_executor/stream.h"
#include "tensorflow/core/framework/device_base.h"
#include "tfrt/concurrency/async_value_ref.h"  // from @tf_runtime

namespace tensorflow {

// XlaHostRecvDeviceContext is a DeviceContext that is intended to be

#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"

namespace tensorflow::quantization {

// Declares pure virtual member functions for a python-side derived class to
// override. This allows calling python implementations from the C++ layer.

absl::Status InitStreamExecutorPlugin(SEInitPluginFn init_fn,
                                      std::string* device_type,
                                      std::string* platform_name);

// This file implements core stream executor base classes in terms of
// the C API defined in stream_executor.h. A class "CSomething" represents a
// "Something" that can be manipulated via calls in the C interface.
class CPlatform : public Platform {

namespace mlir {
namespace TFL {

// A config that controls which passes get run as part TFLite converter.
struct PassConfig {
  explicit PassConfig(quant::QuantizationSpecs specs)
      : quant_specs(std::move(specs)) {}

  // If `emit_builtin_tflite_ops` is true, TF Lite legalization passes will be
  // added, which produces TF Lite ops.
  bool emit_builtin_tflite_ops = true;

#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "tensorflow/core/common_runtime/function_optimization_registry.h"
#include "tensorflow/core/common_runtime/optimization_registry.h"

namespace tensorflow {

// -------------------------------------------------------------------------- //

// executable.
//
// Returns a non-OK Status if no unambiguous choice of device exists.
//
// We choose the device using the following rules:
//
//  - It is an error for `device_names` to contain more than one device of the
//    same type.
//  - GPU is preferred over CPU.
//  - If `allow_mixing_unknown_and_cpu` is true then unknown devices are
//    preferred over CPU.