struct XlaClusterFormationPass
    : public impl::XlaClusterFormationPassBase<XlaClusterFormationPass> {
  void runOnOperation() override;
};

void EncapsulatePartitionedCall(Operation *call_op,
                                mlir::StringAttr callee_name) {
  OpBuilder builder(call_op);
  auto cluster = builder.create<mlir::tf_device::ClusterOp>(