// Adds QuantizationUnitLoc to quantizable layers.
class AddQuantizationUnitLocPass
    : public PassWrapper<AddQuantizationUnitLocPass,
                         OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AddQuantizationUnitLocPass)
  explicit AddQuantizationUnitLocPass() = default;

  StringRef getArgument() const final {

  if (options.enable_inliner) {
    pm.addPass(createInlinerPass());
  }
  pm.addPass(createSymbolDCEPass());
  pm.addNestedPass<func::FuncOp>(CreateTFOptimizePass());
  pm.addNestedPass<func::FuncOp>(createCSEPass());
}

std::unique_ptr<OperationPass<func::FuncOp>> CreateTFOptimizePass() {
  return std::make_unique<TensorFlowOptimizePass>();
}

void RegisterTFOptimizePassPipeline() {

namespace {

// This pass is only available in the tf-opt binary for testing.
class MarkInitializedVariablesTestPass
    : public PassWrapper<MarkInitializedVariablesTestPass,
                         OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MarkInitializedVariablesTestPass)

  StringRef getArgument() const final {
    return "tf-saved-model-mark-initialized-variables-test";
  }

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir::quant {
namespace {

// Applies optimization after quantization.
class OptimizePass
    : public PassWrapper<OptimizePass, OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OptimizePass)

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in

  // Return if the function contains only non-side-effecting ops or TensorArray
  // ops.
  bool HasAtMostTensorArrayEffect(mlir::func::FuncOp func_op) const {
    return set_.count(func_op) > 0;
  }

 private:
  llvm::DenseSet<mlir::func::FuncOp> set_;
};

}  // namespace tfrt_compiler
}  // namespace tensorflow

                             StringRef(content.data(), content.size()));
}

inline bool IsJaxRandomUniform(mlir::func::FuncOp func) {
  return func.getName().contains("tfl_wrapped_jax_random_uniform");
}

inline bool IsJaxRandomNormal(mlir::func::FuncOp func) {
  return func.getName().contains("tfl_wrapped_jax_random_normal");
}

void LegalizeJaxRandomPass::runOnOperation() {

  if (!interface) return false;
  return interface.onlyHasEffect<MemoryEffects::Read>();
}

// Get read only variable handles in `func`.
llvm::DenseSet<ResourceHandle> GetReadOnlyVariables(func::FuncOp func) {
  llvm::DenseMap<ResourceHandle, llvm::SmallVector<Operation *, 4>> resources;

  // Get all VarHandleOps in the function.
  func.walk([&](Operation *op) {
    if (llvm::isa<VarHandleOp>(op)) {

  void runOnOperation() override;
};

void InsertCallOnceOpFromSessionInitializerPass::runOnOperation() {
  ModuleOp module = getOperation();

  for (func::FuncOp init_func_op :
       tf_saved_model::GetInitializerFunctions(module)) {
    for (auto func : module.getOps<func::FuncOp>()) {
      auto dict_attr =
          func->getAttrOfType<mlir::DictionaryAttr>("tf.entry_function");
      if (!dict_attr) continue;

        operands.Insert(op->getOperand(0), ValueConstraint::kShape);
    }

    return success();
  }
};

void TestClusteringPolicyPass::runOnOperation() {
  func::FuncOp func = getOperation();
  ValuesConstraintSet constraints;

  ClusteringPolicySet policies;
  policies.Add<TestOpsClusteringPolicy>();

  // Initialize constraints based on the return type attributes.

    } else if (isa<IdentityNOp, IdentityOp>(op)) {
      value = op->getOperand(res_index);
    } else if (auto call = dyn_cast<CallOpInterface>(op)) {
      func::FuncOp func = dyn_cast<func::FuncOp>(
          call.resolveCallable(&symbol_table_collection_));
      if (!func) break;
      // Check if the function being called has been analyzed. if not,
      // we cannot backtrack the value further.