#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"

namespace mlir::quant::stablehlo {
namespace {

using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::mlir::tf_saved_model::kTfSavedModelInitializerInitType;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;
using ::stablehlo::quantization::QuantizationConfig;
using ::stablehlo::quantization::io::GetLocalTmpFileName;

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_passes.h"

#define DEBUG_TYPE "freeze-global-tensor"

namespace mlir {
namespace tf_saved_model {

namespace {

#define GEN_PASS_DEF_FREEZEGLOBALTENSORSPASS

    if (func::FuncOp func =
            dyn_cast<func::FuncOp>(barg.getOwner()->getParentOp())) {
      SymbolTable symbol_table(func->getParentOfType<ModuleOp>());
      auto global_tensor = tf_saved_model::LookupBoundInputOfType<
          tf_saved_model::GlobalTensorOp>(func, barg.getArgNumber(),
                                          symbol_table);
      ResourceConstructingOps result(global_tensor);
      return result;
    }

            representative_datasets=[
                qc.RepresentativeDatasetConfig(
                    tf_record=qc.TfRecordFile(path=dataset_path)
                )
            ]
        ),
        tf_saved_model=qc.TfSavedModelConfig(tags=[tag_constants.SERVING]),
        pipeline_config=qc.PipelineConfig(
            merge_fusion_with_dequantize=merge_fusion_with_dequantize
        ),
    )

// RUN: stablehlo-quant-opt %s -stablehlo-merge-fusion-with-dequantize -split-input-file -verify-diagnostics | FileCheck %s

// Merge fusion with dequantize for relu case.

module attributes {tf_saved_model.semantics} {
  // CHECK-LABEL: func.func private @merge_relu_fusion
  func.func private @merge_relu_fusion(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {

#include "tsl/platform/status.h"

namespace tensorflow {
namespace quantization {
namespace {

using ::mlir::func::FuncOp;
using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;

// Adds the tensor that initializes the variable through the provided
// `assign_var_op` to the `bundle_writer` for saving to checkpoint. Returns the

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/translate/import_model.h"
#include "tensorflow/core/platform/macros.h"

namespace mlir {
namespace quant {
namespace {

using ::mlir::tf_saved_model::kTfSavedModelExportedNamesAttr;
using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
using ::tensorflow::kImportModelDefaultGraphFuncName;

#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/framework/resource_var.h"

namespace mlir {
namespace tf_saved_model {

std::string GetVariableName(TF::VarHandleOp var_handle_op) {
  // In some cases the shared_name attribute doesn't have the same
  // tensor name in the model, so we first try to use the location

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_executor.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/core/platform/test.h"

namespace mlir::quant {

using ::testing::Test;

class QuantizationTestBase : public Test {
 protected:
  QuantizationTestBase()

#include <functional>
#include <memory>

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"