PackedTensor

class lightning_ir.retrieve.base.packed_tensor.PackedTensor(*args, lengths: Sequence[int] | None = None, **kwargs)[source]

Bases: Tensor

A tensor that contains a sequence of tensors with varying lengths.

__init__(*args, lengths: Sequence[int] | None = None, **kwargs) → None[source]

Initialize the PackedTensor instance.

Parameters:

lengths (Sequence[int] | None) – A sequence of lengths for each segment in the tensor. If provided, the tensor must be created with a total length equal to the sum of these lengths. Defaults to None.

Raises:

  • ValueError – If lengths is None.

  • ValueError – If the sum of lengths does not equal the length of the tensor.

Methods

__init__(*args[, lengths])

Initialize the PackedTensor instance.

lookup(packed_idcs, idcs_lengths[, unique])

Lookup segments in the packed tensor based on provided indices.

to_padded_tensor([pad_value])

Convert the packed tensor to a padded tensor.

Attributes

segmented_tensor

Get the segmented tensor, which is a tuple of tensors split according to the specified lengths.
lookup(packed_idcs: Tensor, idcs_lengths: Sequence[int] | int, unique: bool = False) → PackedTensor[source]

Lookup segments in the packed tensor based on provided indices.

Parameters:

  • packed_idcs (torch.Tensor) – A tensor containing indices to lookup in the packed tensor.

  • idcs_lengths (Sequence[int] | int) – Lengths of the indices for each segment. If a single integer is provided, it is assumed that all segments have the same length.

  • unique (bool) – If True, returns only unique values from the segments. Defaults to False.

Returns:

A new PackedTensor containing the concatenated segments corresponding to the provided indices.

Return type:

PackedTensor

property segmented_tensor: Tuple[Tensor, ...]

Get the segmented tensor, which is a tuple of tensors split according to the specified lengths.

Returns:

A tuple of tensors, each corresponding to a segment defined by the lengths.

Return type:

Tuple[torch.Tensor, …]

to_padded_tensor(pad_value: int = 0) → Tensor[source]

Convert the packed tensor to a padded tensor.

Parameters:

pad_value (int) – The value to use for padding. Defaults to 0.

Returns:

A padded tensor where each segment is padded to the length of the longest segment

in the packed tensor.

Return type:

torch.Tensor