mlir::cast<FlatSymbolRefAttr>(sym_ref).getValue());

    if (!init_func_op)
      return session_initializer.emitOpError()
             << "the initializer function does not exist";

    if (!init_func_op.getFunctionType().getResults().empty())
      return session_initializer.emitOpError()
             << "the initializer function should have no output";

    auto exported_names = GetExportedNames(init_func_op);

    // TODO(aminim) figure out the location story here
    module_ = ModuleOp::create(builder_.getUnknownLoc());
    func_ = func::FuncOp::create(
        builder_.getUnknownLoc(), name,
        builder_.getFunctionType(std::nullopt, std::nullopt));
    module_->push_back(func_);
    builder_ = OpBuilder::atBlockBegin(func_.addEntryBlock());
  }

  void Release() override { delete this; }

  }

  SmallVector<Type> output_types;
  for (const Value output : xla_call_module_op.getOutput()) {
    output_types.push_back(output.getType());
  }

  entry_func_op.setFunctionType(
      rewriter.getFunctionType(arg_types, output_types));

  // Replace argument types and locs.
  Block& entry = entry_func_op->getRegion(0).front();
  for (auto [arg, arg_type, arg_loc] :

  builder.setInsertionPointToEnd(&module.getBodyRegion().back());
  func::FuncOp func_op = builder.create<func::FuncOp>(
      module.getLoc(), name,
      builder.getFunctionType(input_types, output_types));
  func_op.setPrivate();

  return func_op;
}

TF::StatefulPartitionedCallOp EncapsulateOpsInFunc(
    OpBuilder& builder, const llvm::SetVector<Operation*>& ops,