Repository: pytorch
Issue: Enhanced torch.chunk and torch.split 
Labels: good first issue,triaged,enhancement

## 🚀 Feature

Add `drop_remainder: bool = True` support to `torch.chunk` and `torch.split`
+ similar in function to `drop_last` in `torch.utils.data.DataLoader`
+ If the length of the dimension is not perfectly divisible, drop the remaining elements from the returned arrays.

Add `redistribute: bool = True` support to `torch.chunk` and `torch.split`
+ Spread the elements evenly across the chunks/splits if the length of the dimension is not perfectly divisible by the number of chunks or split size.
+ The length of returned tensors differs by at most one element.

## Motivation

Often it is desirable to have evenly sizes arrays, or approximately even sized arrays.

A combination of `drop_remainder` and `redistribute` would enable new intuitive ways to use `torch.split` and `torch.chunk`.

## Pitch

**redistribute**

Add the parameter `redistribute` to `torch.chunk` and `torch.split` which spreads the remainder evenly across the returned tensors. For example:

```python3
# CURRENT:
torch.arange(7).chunk(3, dim=-1, redistribute=False)
>>> (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6]))

# NEW:
torch.arange(7).chunk(3, dim=-1, redistribute=True)
>>> (tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6]))
```

```python3
# CURRENT:
torch.arange(7).split(3, dim=-1, redistribute=False)
>>> (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6]))

# NEW:
torch.arange(7).split(3, dim=-1, redistribute=True)
>>> (tensor([0, 1, 2]), tensor([3, 4]), tensor([5, 6]))
```

**drop_remainder**

Add the parameter `drop_remainder` to `torch.chunk` and `torch.split` which excludes the remainder elements from being returned, and so all the returned tensors will be the same size. For example:

```python3
# CURRENT:
torch.arange(7).chunk(3, dim=-1, drop_remainder=False)
>>> (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6]))

# NEW:
torch.arange(7).chunk(3, dim=-1, drop_remainder=True)
>>> (tensor([0, 1]), tensor([2, 3]), tensor([4, 5]))
```

```python3
# CURRENT:
torch.arange(7).split(3, dim=-1, drop_remainder=False)
>>> (tensor([0, 1, 2]), tensor([3, 4, 5]), tensor([6]))

# NEW:
torch.arange(7).split(3, dim=-1, drop_remainder=True)
>>> (tensor([0, 1, 2]), tensor([3, 4, 5]))
```

**drop_remainder & redistribute interaction**

`drop_remainder` should be applied before `redistribute`:
- thus if both `drop_remainder=True`, then there is never a need to redistribute elements so in this case `redistribute` can be ignored.

## Alternatives

New functions that provide similar behavior or act over the outputs from `chunk` or `split`, however this might be confusing in the first case or inefficient in the second case.

Possible synonyms for `drop_remainder` include:
+ `drop_last` matching `torch.utils.data.DataLoader`, but it may be argued that does not convey the correct behaviour in all cases?

Possible synonyms for `redistribute` include
+ `redistribute_remainder`
+ `distribute_remainder` 
+ `spread_remainder`
+ `distribute`
+ `spread`

The `_remainder` does not describe the movement of the remainder/last elements well, as internal elements are shifted too.

## Additional context

N/A