#include "xla/client/local_client.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/local_device.h"
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"

#include "tensorflow/compiler/tf2xla/xla_argument.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/xla_computation.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"

namespace tensorflow {

// Lowers MLIR module to XLA HLO inside an XlaComputation. The input module

#include "stablehlo/dialect/VhloOps.h"  // from @stablehlo
#include "stablehlo/dialect/VhloTypes.h"  // from @stablehlo
#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"

namespace mlir {

// duplicated from
// https://github.com/openxla/stablehlo/blob/e5ad51715a11721c78b6748ab5de7945df24b1b8/stablehlo/transforms/StablehloLegalizeToVhlo.cpp#L756

    CAPSULE = 10,
  };
  TaggedValue() : type_(NONE), data_() {}

  /// Move assignment operator.
  TaggedValue& operator=(TaggedValue&& v) {
    destroy();
    MoveIntoUnion(std::move(v));
    return *this;
  }
  /// Move constructor.
  TaggedValue(TaggedValue&& v) : type_(NONE) { MoveIntoUnion(std::move(v)); }
  /// Copy constructor.
  TaggedValue(const TaggedValue& v) : type_(NONE) { CopyIntoUnion(v); }

#include "tensorflow/compiler/mlir/quantization/common/quantization_lib/quantization_utils.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/ops/stablehlo_op_quant_spec.h"
#include "tensorflow/core/framework/types.pb.h"

namespace mlir::quant::stablehlo {

// Checks whether an op is connected with a quantized composite function. If
// not, the same-scale op will not be quantized. This decision is based on the