struct FreezeGlobalTensorsPass
    : public impl::FreezeGlobalTensorsPassBase<FreezeGlobalTensorsPass> {
  explicit FreezeGlobalTensorsPass(bool allow_mutable_tensors) {
    this->allow_mutable_tensors = allow_mutable_tensors;
  }
  void runOnOperation() override;
};

void FreezeGlobalTensorsPass::runOnOperation() {
  auto module = getOperation();

// RUN: tf-opt "-tf-saved-model-freeze-global-tensors=allow-mutable-tensors=true" -split-input-file %s | FileCheck %s

module attributes {tf_saved_model.semantics} {

  // Test case: Do not fail if the tensor is mutable but allow_mutable_tensors is true

  "tf_saved_model.global_tensor"() { is_mutable, sym_name = "v", type = tensor<f32>, value = dense<1.0> : tensor<f32> } : () -> ()

std::unique_ptr<OperationPass<ModuleOp>> CreateOptimizeGlobalTensorsPass();

// Creates a pass that freezes tf_saved_model.global_tensor ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateFreezeGlobalTensorsPass(
    bool allow_mutable_tensors = false);

// Creates a pass that freezes tf_saved_model.asset ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateFreezeAssetsPass(
    std::string saved_model_dir = "");

  // freezing of TF v2 GlobalTensor ops. (Ref: b/206855389)
  pm.addPass(mlir::tf_saved_model::CreateOptimizeGlobalTensorsPass());
  pm.addPass(mlir::tf_saved_model::CreateFreezeGlobalTensorsPass(
      /*allow_mutable_tensors=*/true));

  // Generic MLIR optimization passes.
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addPass(mlir::TF::CreateTFShapeInferencePass());

  // Legalizes TF UniformQuantized types into MHLO.

  // freezing of TF v2 GlobalTensor ops. (Ref: b/206855389)
  pm.addPass(mlir::tf_saved_model::CreateOptimizeGlobalTensorsPass());
  pm.addPass(mlir::tf_saved_model::CreateFreezeGlobalTensorsPass(
      /*allow_mutable_tensors=*/true));

  // Generic MLIR optimization passes.
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addPass(mlir::TF::CreateTFShapeInferencePass(input_arg_shapes));

    pass lowers from saved model to pure TF. Hence it fails, if it cannot lower.
  }];

  let constructor = "::mlir::tf_saved_model::CreateFreezeGlobalTensorsPass()";

  let options = [
    Option<"allow_mutable_tensors", "allow-mutable-tensors", "bool", /*default=*/"false",
           "Allows mutable tensors to be in the graph.">
  ];
}

def OptimizeGlobalTensorsPass : Pass<"tf-saved-model-optimize-global-tensors", "ModuleOp"> {

    // This pass 'freezes' immutable global tensors and inlines them as tf
    // constant ops.
    pass_manager->addPass(mlir::tf_saved_model::CreateFreezeGlobalTensorsPass(
        /*allow_mutable_tensors=*/pass_config.enable_tflite_variables));
  }

  if (!saved_model_dir.empty()) {
    // This pass 'freezes' tf saved model asset ops and inlines as string values
    // in a format of the tf constant op.