indices_or_sections parameter of tensor_split() doesn't work while sections, indices or tensor_indices_or_sections parameter works

[hyperkai]

📚 The doc issue

The doc of tensor_split() says that there is indices_or_sections parameter as shown below but there isn't:

indices_or_sections (Tensor, int or list or tuple of ints) –

Actually, indices_or_sections parameter doesn't work while sections(int), indices(tuple of int or list of int) or tensor_indices_or_sections(tensor of int) parameter works as shown below:

import torch

my_tensor = torch.tensor([[0, 1, 2, 3],
                          [4, 5, 6, 7],
                          [8, 9, 10, 11]])
torch.tensor_split(input=my_tensor, indices_or_sections=1) # Error
torch.tensor_split(input=my_tensor, indices_or_sections=(1, 2)) # Error
torch.tensor_split(input=my_tensor, indices_or_sections=torch.tensor(1)) # Error
torch.tensor_split(input=my_tensor, indices_or_sections=torch.tensor([1, 2])) # Error

torch.tensor_split(input=my_tensor, sections=1)
# (tensor([[0, 1, 2, 3],
#          [4, 5, 6, 7],
#          [8, 9, 10, 11]]),)

torch.tensor_split(input=my_tensor, indices=(1, 2))
# (tensor([[0, 1, 2, 3]]),
#  tensor([[4, 5, 6, 7]]),
#  tensor([[8, 9, 10, 11]]))

torch.tensor_split(input=my_tensor, tensor_indices_or_sections=torch.tensor(1))
# (tensor([[0, 1, 2, 3],
#          [4, 5, 6, 7],
#          [8, 9, 10, 11]]),)

torch.tensor_split(input=my_tensor, tensor_indices_or_sections=torch.tensor([1, 2]))
# (tensor([[0, 1, 2, 3]]),
#  tensor([[4, 5, 6, 7]]),
#  tensor([[8, 9, 10, 11]]))

Suggest a potential alternative/fix

So, the doc of tensor_split() should be as shown below:

indices(tuple of int or list of int), sections(int) or tensor_indices_or_sections(tensor of int) –

cc @albanD