---

FoldedTensor: PyTorch extension for handling deeply nested sequences of variable length

foldedtensor is a PyTorch extension that provides efficient handling of tensors containing deeply nested sequences variable sizes. It enables the flattening/unflattening (or unfolding/folding) of data dimensions based on a inner structure of sequence lengths. This library is particularly useful when working with data that can be split in different ways and enables you to avoid choosing a fixed representation.

Installation

The library can be installed with pip:

pip install foldedtensor

Features

• Support for arbitrary numbers of nested dimensions
• No computational overhead when dealing with already padded tensors
• Dynamic re-padding (or refolding) of data based on stored inner lengths
• Automatic mask generation and updating whenever the tensor is refolded
• C++ optimized code for fast data loading from Python lists and refolding
• Flexibility in data representation, making it easy to switch between different layouts when needed

Examples

At its simplest, foldedtensor can be used to convert nested Python lists into a PyTorch tensor:

from foldedtensor import as_folded_tensor

ft = as_folded_tensor(
    [
        [0, 1, 2],
        [3],
    ],
)
# FoldedTensor([[0, 1, 2],
#               [3, 0, 0]])

You can also specify names and flattened/unflattened dimensions at the time of creation:

import torch
from foldedtensor import as_folded_tensor

# Creating a folded tensor from a nested list
# There are 2 samples, the first with 5 lines, the second with 1 line.
# Each line contain between 1 and 2 words.
ft = as_folded_tensor(
    [
        [[1], [], [], [], [2, 3]],
        [[4, 3]],
    ],
    data_dims=("samples", "words"),
    full_names=("samples", "lines", "words"),
    dtype=torch.long,
)
print(ft)
# FoldedTensor([[1, 2, 3],
#               [4, 3, 0]])

Once created, you can change the shape of the tensor by refolding it:

# Refold on the lines and words dims (flatten the samples dim)
print(ft.refold(("lines", "words")))
# FoldedTensor([[1, 0],
#               [0, 0],
#               [0, 0],
#               [0, 0],
#               [2, 3],
#               [4, 3]])

# Refold on the words dim only: flatten everything
print(ft.refold(("words",)))
# FoldedTensor([1, 2, 3, 4, 3])

The tensor can be further used with standard PyTorch operations:

# Working with PyTorch operations
embedder = torch.nn.Embedding(10, 16)
embedding = embedder(ft.refold(("words",)))
print(embedding.shape)
# torch.Size([5, 16]) # 5 words total, 16 dims

refolded_embedding = embedding.refold(("samples", "words"))
print(refolded_embedding.shape)
# torch.Size([2, 5, 16]) # 2 samples, 5 words max, 16 dims

Pooling spans

lengths.make_indices_ranges maps half open spans to storage positions, excluding padding even when a span crosses rows. It returns the expanded indices, start offsets and the span id of each selected item.

import torch
import foldedtensor as ft

tensor = ft.as_folded_tensor([[1.0, 2.0], [3.0]], full_names=("sample", "word"))
indices, offsets, spans = tensor.lengths.make_indices_ranges(
    begins=(torch.tensor([0, 1]),),
    ends=(torch.tensor([2, 3]),),
    indice_dims=("word",),
)
pooled = torch.nn.functional.embedding_bag(
    indices,
    tensor.as_tensor().reshape(-1, 1),
    offsets,
    mode="mean",
)
assert pooled.tolist() == [[1.5], [2.5]]

Boundary mapping and expansion run in C++, using prefix offsets from the sequence lengths and the refolding indexer for padded layouts. Range indices are computed on CPU and returned on the input device. Embedding gathering and pooling run on the embedding tensor's device.

Benchmarks

View the comparisons of foldedtensor against various alternatives here: docs/benchmarks.

Comparison with alternatives

Unlike other ragged or nested tensor implementations, a FoldedTensor does not enforce a specific structure on the nested data, and does not require padding all dimensions. This provides the user with greater flexibility when working with data that can be arranged in multiple ways depending on the data transformation. Moreover, the C++ optimization ensures high performance, making it ideal for handling deeply nested tensors efficiently.

Here is a comparison with other common implementations for handling nested sequences of variable length:

Feature	NestedTensor	MaskedTensor	FoldedTensor
Inner data structure	Flat	Padded	Arbitrary
Max nesting level	1	1	∞
From nested python lists	No	No	Yes
Layout conversion	To padded	No	Any
Reduction ops w/o padding	Yes	No	No