scikit_learn train_test_splitを用いて要素数に偏りがないようにテストデータを作成したい。

前提・実現したいこと

クラスごとのラベルの数に偏りのある学習データをscikit-learnのtrain_test_splitを用いて学習データと訓練データに分割したいと考えています。

発生している問題

クラスごとに、分割する割合を決めたいのですがうまく実装できません。

クラスによってデータを集めることが難しいものを学習データとあるので、どうしても大きな偏りができてしまいます。
そのため、テストデータのクラスごとの要素数が極端に少なくならないようにする方法を探しています。

該当のソースコード

python
1#学習データを想定(1000*2*2*1の配列)
2a = np.arange(4000).reshape(1000,2,2,1)
3
4#学習データのラベルを想定(1000*1の配列)
5b = np.array([0] * 500 + [1] *200 + [2] * 200 + [3] * 100).reshape(1000,1)
6
7#分割
8X_train, X_test, Y_train, Y_test = train_test_split(a, b,test_size=0.1,random_state=1,stratify=b)
9
10#学習データとテストデータの次元と要素数の確認
11print(X_train.shape)
12print(Y_train.shape)
13print(X_test.shape)
14print(Y_test.shape)
15
16#学習データとテストデータのクラスの要素数の確認
17print((Y_test == 0).sum())
18print((Y_test == 1).sum())
19print((Y_test == 2).sum())
20print((Y_test == 3).sum())
21
22#実行結果
23(750, 2, 2, 1)
24(750, 1)
25(250, 2, 2, 1)
26(250, 1)
27125
2850
2950
3025
31

入力データは私が実際に用いるデータの次元を想定し、4次元にしました。

この場合、クラスラベル3の要素数が少ないので、クラス3をもう少し学習データに回したり、テストデータ側に回したり、できればと考えています（例：テストデータ内でクラス3の要素数が25、と少ないので、クラス0 1 2の要素数はそのままに、クラス3だけ50に増やす）。
または、テストデータ内でクラス3の要素数はそのままにして、クラス0 1 2を学習データ側に回して、テストデータ内の各クラスの要素数を統一したいと考えています(例：テストデータ内の要素数を全て50に統一する)。