mnistデータセットのような二次元配列のデータセットの作り方

実現したいこと

・１００行×１００列の数字の配列のデータが１００００個ほどあり、これから全て別々のｃｓｖファイルに保存しようと考えています。データセットが完成したらCNNで学習させます。

・アヤメの分類を行った際は一つのｃｓｖファイルの同じ行に学習に使う４つの要素(x)とラベル(y)を並べておき(150行×５列のｃｓｖファイル)、np.splitでx,yに分け、train_test_splitでxtrain,ytrain,xtest,ytestに分けるという方法を取っていましたが、今回は学習に使う要素が二次元であるため同じ方法が取れずにどうすれば良いか分からない、という状況です。

・mnistデータセットのxtrain,ytrainのように、データとラベルを対応させたデータセットを作りたいと考えています。ただし、交差検証をすることを見据えて、mnistデータセットの通りxtrain,ytrain,xtest,ytestの４つに分けるという訳ではなく、x,yの二つに分けたいです。

解決方法をご存じの方がいらっしゃればご教授いただければ幸いです。よろしくお願いいたします。