#ifndef TENSORFLOW_COMPILER_JIT_XLA_COMPILE_UTIL_H_
#define TENSORFLOW_COMPILER_JIT_XLA_COMPILE_UTIL_H_

#include <memory>
#include <string>

#include "tensorflow/compiler/tf2xla/xla_argument.h"
#include "tensorflow/core/graph/graph.h"

namespace tensorflow {
// The number of compiler threads to use for asynchronous device compilation.
inline constexpr int64_t kNumAsyncDeviceCompilerThreads = 10;

enum class DeviceCompileMode {

#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V1_TF_DIALECT_TO_EXECUTOR_H_

#include "absl/base/attributes.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "tensorflow/core/platform/status.h"

namespace tensorflow {
namespace tf2xla {
namespace v1 {

// Given the input Module op that's in the Tensorflow Dialect, convert the MLIR

#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 mlir::quant::stablehlo {

// Performs int8 weight-only quantization on dot_general ops.
//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in

 private:
  // Instrumentation hooks. These hooks don't need to be thread-safe. The pass
  // manager runs each pass for the entire module, then it walks through
  // each op in the module and runs the pass on them, may be in async mode.
  void runBeforePass(Pass *pass, Operation *module) override;
  void runAfterPass(Pass *pass, Operation *module) override;
  void runAfterPassFailed(Pass *pass, Operation *module) override;

#define TENSORFLOW_COMPILER_JIT_XLA_HOST_SEND_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 {

// XlaHostSendDeviceContext is a DeviceContext that is intended to be

#include "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/core/platform/types.h"

namespace tensorflow {
namespace generator {
namespace cpp {

class Renderer {
 public:
  explicit Renderer(RendererContext context);

 protected:
  // Append a blank line.

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

namespace mlir::quant::stablehlo {

// Component for static-range post-training quantization (PTQ).
// TODO: b/320607042 - Add tests in python level.

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

namespace mlir::quant::stablehlo {

// Performs post-calibration graph transformation as part of post-training
// static-range quantization.
//

namespace mlir {
namespace odml {

// Given a Composite op that wraps a core.aten.avg_pool2d, returns the padding
// configuration required for the `tfl.pad` if the padding part of the op is
// to be done before average pooling.
DenseIntElementsAttr GetPadOpAttr(Builder& builder, mhlo::CompositeOp op);

// Given a Composite op that wraps a core.aten.avg_pool2d, and assuming that
// the padding part is extracted into a tfl.pad op prior to a

#include "tensorflow/core/framework/function.pb.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace tfr {

extern const char* const kTFRLibEnv;

using tsl::StatusOr;

// An wrapper for all the objects used to decompose a module (graph mode) and
// node_def (eager mode). Note that this class owns the decomposition library.
class TFRDecomposeContext {