# Set `force_graph_mode_calibration` to True to avoid skipping op execution,
  # which are not connected to return ops, during calibration execution.
  # TODO: b/335031954 - Bring back support to run calibration in Eager mode.
  logging.debug(
      'Setting `force_graph_mode_calibration = True` to ensure the calibration'
      ' mode is executed properly.'
  )
  quantization_options.force_graph_mode_calibration = True

#include "mlir/Support/LLVM.h"  // from @llvm-project
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/component.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/calibration/statistics.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/context.h"

// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-insert-calibration-statistics-saver='aggregator-ops-to-ignore=skipping_id' | FileCheck %s

func.func @serving_default(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x2x2x2xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}} {

    )

    # Run quantization the first time, calibration is expected to be run.
    with self.assertLogs(level='INFO') as info_logs:
      quantization.quantize_saved_model(
          self._input_saved_model_path,
          self._output_saved_model_path,
          config,
      )
      self.assertTrue(
          self._any_log_contains(
              'Calibration step is executed in graph mode.',

        # Restore the logger verbosity.
        logging.set_verbosity(prev_log_level)

      self.assertNotEmpty(info_logs.records)
      self.assertTrue(
          self._any_log_contains(
              'Calibration step is executed in graph mode.',
              info_logs.records,
          )
      )


class TensorNamePreservationTest(quantize_model_test_base.QuantizedModelTest):

                   FlatSymbolRefAttr::get(builder.getStringAttr(func_name)));
  // Store the custom attribute to restore the function name when loading it
  // back in the post calibration stage. As mentioned above, the above entry
  // function attribute is not reliable.
  call_op->setAttr(kOriginalStablehloEntryFunctionAttrName,
                   builder.getStringAttr(func_name));

    input_model_ = ReadModel(internal::kConvModelWith0Plus10Weights);
    readonly_model_ = input_model_->GetModel();
    model_ = UnPackFlatBufferModel(*readonly_model_);
    // Flatbuffer is missing calibration data -- add dummy params.
    auto& subgraph = model_.subgraphs[0];
    auto* input = subgraph->tensors[subgraph->inputs[0]].get();
    auto* output = subgraph->tensors[subgraph->outputs[0]].get();

#include "tsl/platform/status.h"
#include "tsl/platform/status_matchers.h"

namespace tensorflow {
namespace {

using ::stablehlo::quantization::CalibrationOptions;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsMap;
using ::testing::Contains;
using ::testing::ElementsAre;
using ::testing::HasSubstr;
using ::testing::Key;
using ::testing::SizeIs;
using ::tsl::testing::StatusIs;

    *   `converter.convert()` no longer returns a `tf.function`. Now the
        function must be accessed from the saved model.
    *   The `converter.calibrate()` method has been removed. To trigger
        calibration, a `calibration_input_fn` should be provided to
        `converter.convert()`.

*   Other:

    *   Fix accidental quadratic graph construction cost in graph-mode
        `tf.gradients()`.