feat: to/from PyTorch JaggedTensor

[maxymnaumchyk]

No description provided.

[codecov]

Codecov Report

Attention: Patch coverage is 23.25581% with 66 lines in your changes missing coverage. Please review.

Project coverage is 81.85%. Comparing base (b749e49) to head (c68321c).
Report is 162 commits behind head on main.

Files with missing lines	Patch %	Lines
src/awkward/operations/ak_to_jaggedtensor.py	20.00%	40 Missing ⚠️
src/awkward/operations/ak_from_jaggedtensor.py	23.52%	26 Missing ⚠️

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/operations/__init__.py	100.00% <100.00%> (ø)
src/awkward/operations/ak_from_jaggedtensor.py	23.52% <23.52%> (ø)
src/awkward/operations/ak_to_jaggedtensor.py	20.00% <20.00%> (ø)

... and 105 files with indirect coverage changes

[maxymnaumchyk]

@jpivarski should I also leave out a "keep_regular" parameter since it does the same as ak.from_regular()? It's kind of the same situation we talked about today (about "padded" parameter).

[jpivarski]

You're right: it is. The situation is that we should be providing an interface to the user that's like Lego bricks that they can put together however they like. If there's an alternative way of doing something, it shouldn't be a feature of the new functions, because then we'd have to explain why someone would use one or the other.

I agree that the padded and keep_regular arguments are more convenient if that's exactly what someone wants; passing padded=True or keep_regular=True is easier than the multi-step process it would be with the other method. However, these shortcuts would only work in the cases that these functions apply to, which are limited by the capabilities of TensorFlow and PyTorch's ragged array implementations.

[maxymnaumchyk]

thanks for such a detailed answer!

[jpivarski]

This is great work! But it might be solving the wrong problem. As we discussed at the meeting, this fbgemm-gpu-cpu module is not something ML users seem to be familiar with, so adding to/from functions wouldn't help them. It doesn't seem to be the interface that they use to implement DeepSets and GNNs, the ML models that might actually involve ragged data. So we're going to follow-up with ML experts to find out what interfaces they really do need.

Meanwhile, as discussed at the meeting, you'll be adding

• ak.to_torch using Content.to_backend_array
• ak.from_torch

for rectilinear arrays only (allow_record=False and allow_missing=False, following ak_to_numpy.py).

Another two functions,

• ak.parents_index with axis #3256
• ak.records_to_regular to convert [{"x": 1, "y": 2}, {"x": 3, "y": 4}] into [[1, 2], [3, 4]] #3257

would be needed to pre-process an idiomatic Awkward Array into the kind of interface that PyTorch-Geometric needs, which isn't one RaggedTensor object like TensorFlow; it's a few, separate, completely rectilinear arrays. "The way to do it" needs to be explained as a User Guide that puts all of these functions together, rather than a single function that tries to do everything in one call.

Although I'm setting this to "request changes," we'll likely be closing this PR and following up with new ones.