For classification you can use the `stratify` parameter:

> **stratify: array-like or None (default=None)** 
>
> If not None, data is split in a stratified fashion, using this as the class labels.

See [sklearn.model_selection.train_test_split][1]. For example:

```
x, x_test, y, y_test = train_test_split(xtrain,labels,test_size=0.2, stratify=labels) 
```

This will ensure the class distribution is similar between train and test data. 
*(side note: I have tossed the `train_size` parameter since it will be automatically determined based on `test_size`)*

For regression there is, to my knowledge, no current implementation in scikit learn. But you can find a discussion and manual implementation [here][2] and [here][2] with regards to cross-validation. 


  [1]: https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html
  [2]: https://github.com/scikit-learn/scikit-learn/issues/4757