#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 "absl/strings/str_cat.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/lib/strings/strcat.h"

namespace tensorflow {

/// @defgroup core Core Tensorflow API

class Output;

/// @addtogroup core
/// @{

/// Represents a node in the computation graph.

    if (llvm::isa<FixedOutputRangeInterface, SameScalesOpInterface>(def) ||
        !def->hasTrait<OpTrait::quant::QuantizableResult>()) {
      return failure();
    }

    // This op should not clobber def, if more than one requant of this value.
    if (!pre_quantized.hasOneUse()) {
      return failure();
    }

    op.emitWarning("Remove trivial `rescale` op. Please fix the source graph.");