}
  }
}

void ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass::
    runOnOperation() {
  ModuleOp module_op = getOperation();

  func::FuncOp main_func = FindMainFuncOp(module_op);
  if (!main_func) return;

  // In case the model has tf.StatefulPartitionedCallOp or tf.PartitionedCallOp,
  // we recursively find called functions and process StableHLO ops in them.