::mlir::quantfork::QuantizationForkDialect>();
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrepareQuantizeDRQPass)

  // Constructor used by the PassRegistration and enforce int8 quantization.
  // This is only used by test.
  explicit PrepareQuantizeDRQPass() : op_set_(OpSet::UNIFORM_QUANTIZED) {
    quant_specs_.inference_type = tensorflow::DT_QINT8;
  }

        failed(ApplyVhloToStablehloPatterns(module)) ||
        failed(ApplyUnrealizedCastCanonicalization(module)))
      return signalPassFailure();
  }
};

static PassRegistration<LegalizeStablehloToVhloPass> pass_s2v;
static PassRegistration<LegalizeVhloToStablehloPass> pass_v2s;

}  // namespace odml

                    ::mlir::quantfork::QuantizationForkDialect>();
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrepareQuantizePass)

  // Constructor used by the PassRegistration and enforce uint8 quantization.
  // This is only used by test.
  explicit PrepareQuantizePass() {
    quant_specs_.inference_type = tensorflow::DT_QINT8;
  }

  // Constructor used by manually creating the pass.

    custom_op.erase();
  });
}

std::unique_ptr<OperationPass<func::FuncOp>> CreateTflToStablehloPass() {
  return std::make_unique<TflToStablehloPass>();
}

static PassRegistration<TflToStablehloPass> pass;

}  // namespace odml

      }
    }
    return llvm::StringRef("");
  };

  return CreateImportQuantStatsPass(get_name_func, stats_str);
}

// Registers this pass with default values, only for test
static PassRegistration<ImportQuantStatsPass> pass([] {
  return CreateImportQuantStatsPassForTFControlDialect(quantize_stats);
});

}  // namespace quant

      return signalPassFailure();
    }
  }
};

std::unique_ptr<Pass> createUnfuseBatchNormPass() {
  return std::make_unique<UnfuseMhloBatchNormPass>();
}

static PassRegistration<UnfuseMhloBatchNormPass> pass;

}  // namespace odml

    }

    op->emitRemark("Not supported quantization data type.");
    return failure();
  }

 protected:
  tensorflow::quantization::QuantizationOptions quant_options_;
};

static PassRegistration<QuantizeWeightsPass> pass;

void QuantizeWeightsPass::runOnOperation() {
  MLIRContext* ctx = &getContext();
  auto module_op = getOperation();
  RewritePatternSet patterns(ctx);

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

  // Constructor used by the PassRegistration and only used by test.
  explicit QuantizePass() {
    quant_specs_.inference_type = tensorflow::DT_QINT8;
  }

  // Constructor used by manually creating the pass.

      return signalPassFailure();
    }
  }
};

std::unique_ptr<Pass> createFoldBroadcastPass() {
  return std::make_unique<FoldBroadcastPass>();
}

static PassRegistration<FoldBroadcastPass> pass;

}  // namespace odml

struct QuantizePass : public impl::QuantizePassBase<QuantizePass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizePass)

  // Constructor used by the PassRegistration and only used by test.
  explicit QuantizePass() { quant_specs.inference_type = tensorflow::DT_QINT8; }

  // Constructor used by manually creating the pass.