#include "mlir/IR/BuiltinOps.h"  // from @llvm-project

namespace tflite {

// Returns serialized string for the generated flatbuffer.
std::optional<std::string> ExportRuntimeMetadata(mlir::ModuleOp module);

}  // namespace tflite

TEST_F(PreCalibrationComponentTest,
       HasCustomAggregatorOpAndQuantizableFuncForSimpleDotGeneral) {
  PreCalibrationComponent component(ctx_.get());
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    module attributes {} {
      func.func @main(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> attributes {} {
        %0 = stablehlo.constant dense<1.0> : tensor<4x3xf32>

          signature_def_map,
      absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases);

  // Runs the static-range post-training quantization (PTQ) on `module_op`.
  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,
      const ::stablehlo::quantization::QuantizationConfig& config) override;

 private:
  // A non-owning `MLIRContext`. This `MLIRContext` should exceed the lifetime
  // of `StaticRangePtqComponent`.

#include "tsl/lib/core/status_test_util.h"
#include "tsl/platform/statusor.h"

namespace mlir::quant::stablehlo {
namespace {

using ::mlir::DialectRegistry;
using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::tensorflow::monitoring::testing::CellReader;
using ::testing::Test;

static constexpr char kMetricsName[] =
    "/tensorflow/core/tf2xla/tf_quant_op_count";

class DecomposeTFOpsPass
    : public PassWrapper<DecomposeTFOpsPass, OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DecomposeTFOpsPass)

  explicit DecomposeTFOpsPass(std::optional<ModuleOp> external_tfr_module)
      : external_tfr_module_(external_tfr_module) {}

  StringRef getArgument() const final { return "tfr-decompose"; }

  StringRef getDescription() const final {

namespace mlir {
namespace tf_saved_model {

// Initializes all variables in Session Init function for all variables in
// 'session'.
LogicalResult InitializeVariablesInSessionInitializer(
    ModuleOp module, tensorflow::Session *session);

}  // namespace tf_saved_model

}  // namespace mlir

include "mlir/Pass/PassBase.td"

// Internal TF2XLA Bridge Passes used during interence only.
def InferenceMetricsPass : Pass<"tf2xla-inference-metrics-pass", "ModuleOp"> {
  let summary = "A pass to go over a Module and collect various metrics";
  let description = [{
    Collects metrics about a Module during inference, such as which ops are used.
  }];

  if (Operation *nested_module = symbol_table.lookup(kNestedModule)) {
    nested_module->emitOpError("unexpected already present outlined module.");
    return signalPassFailure();
  }
  ModuleOp outlined_module = ModuleOp::create(getOperation().getLoc());
  outlined_module->setAttrs(getOperation()->getAttrDictionary());
  outlined_module->setAttr(SymbolTable::getSymbolAttrName(),

using ::tensorflow::quantization::RunPasses;

WeightOnlyPtqComponent::WeightOnlyPtqComponent(absl::Nonnull<MLIRContext*> ctx)
    : ctx_(ABSL_DIE_IF_NULL(ctx)) {}  // Crash OK

absl::StatusOr<ModuleOp> WeightOnlyPtqComponent::Run(
    ModuleOp module_op, const QuantizationConfig& config) {
  TF_RETURN_IF_ERROR(RunPasses(
      kName, /*add_passes_func=*/
      [&config](PassManager& pm) {

namespace mlir {
namespace tf_saved_model {

// Creates GlobalTensorOp for each variable from function arguments and converts
// them to the corresponding saved model arguments.
LogicalResult LiftVariables(ModuleOp module, ::tensorflow::Session* session);

}  // namespace tf_saved_model
}  // namespace mlir