Request for torch.split to accept a tensor for input split_size_or_sections

[edqwerty10]

🚀 Feature

For the pytorch operator torch.split(tensor, split_size_or_sections, dim) we would like the input split_size_or_sections to also handle tensors. Right now it handles list of ints and an int.

Motivation

When tracing a model, we currently do split_size_or_sections=tensor.tolist() when calling torch.split but tracing can't record this conversion of tensor to list of ints so the traced model fails on other inputs. This is currently blocking a diff that tries to improve runtime performance by using tensor operations, D24595761.

Pitch

Have torch.split(tensor, tensor, int) work so that tracing can properly record the operations for a traced model.

Alternatives

We currently don't have alternatives (I think). We are migrating to a scripted model (which works in this case) but support for tracing is still requested.

Additional context

Here is a diff D24595761 that tries to improve performance by using torch operators only but fails for a traced model

[dzhulgakov]

The issue is that there's some shady conversion going on here:

pytorch/torch/tensor.py

Line 509 in 3549141

return super(Tensor, self).split_with_sizes(split_size, dim)

so today it seems that we convert 1-element tensor as a "number of split" not "sizes of splits".

What requested can be safely applied to split_with_sizes though because it's unambiguous.

cc @gchanan @mruberry for opinions

[mruberry]

We also have a NumPy compatible tensor_split that should throw an error currently when passed a tensor. We could avoid a BC concern and document this function treating a CPU tensor as a list.

This raises the more general questions:

• do we really want CPU tensors to be interpreted as lists consistently?
• is tracing going to continue to be around, or do we expect changes to better-support tracing are temporary?

I think it'd be OK (from a UX perspective) to consistently interpret CPU tensors as lists when an operand can take a list and as a scalar when an operand can only be a scalar. Device types like XLA will probably want the tensor to be passed as an XLA tensor, however. See #31558 and cc @ailzhang. Unfortunately I don't think we can always match these tensors to the device type that will run the operation because passing a CUDA tensor and converting it to a list would cause cross-device data movement.

For the second question, I suppose tracing will be around long enough that adding this support to a few functions would be OK. Especially since if we used tensor_split I don't think we'd have any BC concerns.

Do you think tensor_split would work for you, @edqwerty10?

[edqwerty10]

Hi @mruberry, this is great thank you for the breakdown. Yes tensor_split should work for us as well!

[mruberry]

OK. We would accept a PR updating tensor_split to accept a CPU tensor in place of a list. Note this is consistent with NumPy:

import numpy as np
a = np.array((1, 2, 3, 4))
np.array_split(a, np.array((1, 2)))
: [array([1]), array([2]), array([3, 4])]

np.array_split(a, [1, 2])
: [array([1]), array([2]), array([3, 4])]

[mruberry]

The fastest way would be for you to submit a PR implementing the behavior. If that's a pain you can ping me internally and we can prioritize the request.

[xw285cornell]

Any chance we can allow a cuda tensor for indices?

[mruberry]

Any chance we can allow a cuda tensor for indices?

No, because the indices tensor is used to define the tensor outputs, and the metadata for tensors lives on the CPU, so the CPU has to be able to access the data in indices.

[mruberry]

Closing since we resolved this request by implementing the functionality in torch.tensor_split, and the request for this functionality in torch.split is a dupe of #16703.