Nested Module API Reference

This page provides the complete API reference for the batchtensor.nested module.

The nested module contains functions for manipulating nested data structures (dictionaries, lists, tuples) containing PyTorch tensors. All operations recursively apply to tensors within the nested structure while preserving the structure.

For usage examples and tutorials, see the Nested Operations User Guide.

batchtensor.nested

Functions for manipulating nested data structures containing PyTorch tensors.

This module provides functions for working with nested data structures (dictionaries, lists, tuples) that contain PyTorch tensors. All functions recursively apply operations to every tensor in the structure while preserving the structure itself.

The module uses the coola library's recursive application capabilities to handle arbitrary nesting levels and mixed data structures. This allows you to work with complex batched data (e.g., a dictionary of lists of tensors) using the same simple API as individual tensors.

All functions in this module follow these conventions

• Functions ending with _along_batch operate on the batch dimension (dim=0)
• Functions ending with _along_seq operate on the sequence dimension (dim=1)
• All tensors in the nested structure must have compatible shapes for the operation being performed
• The output structure mirrors the input structure
• The data parameters and return values are annotated with batchtensor.nested.types.NestedTensor, a recursive type alias for "a tensor, or a mapping/sequence of NestedTensor"

Function Categories

Reduction operations: Aggregate values along a dimension - sum_along_batch, mean_along_batch, max_along_batch, etc.

Slicing operations: Extract subsets of data - slice_along_batch, select_along_batch, chunk_along_batch - slice_along_seq, select_along_seq, chunk_along_seq

Joining operations: Combine multiple nested structures - cat_along_batch, cat_along_seq, repeat_along_seq - stack_along_batch, stack_along_seq

Padding operations: Pad ragged batches - pad_along_batch, pad_along_seq

Permutation operations: Reorder elements - permute_along_batch, shuffle_along_batch - permute_along_seq, shuffle_along_seq

Indexing operations: Select specific elements - index_select_along_batch, index_select_along_seq

Comparison operations: Sort and find extrema - sort_along_batch, argsort_along_batch

Mathematical operations: Element-wise and cumulative operations - cumsum_along_batch, cumprod_along_batch - abs, exp, log, clamp

Trigonometric operations: Standard trigonometric functions - sin, cos, tan, asin, acos, atan - sinh, cosh, tanh, asinh, acosh, atanh

Conversion operations: Convert between tensor types and formats - as_tensor, from_numpy, to_numpy, to

Example

Working with nested dictionaries:

>>> import torch
>>> from batchtensor import nested as bn
>>> # Create a nested batch structure
>>> batch = {
...     "inputs": torch.tensor([[1, 2], [3, 4], [5, 6]]),
...     "labels": torch.tensor([0, 1, 0]),
...     "metadata": {"ids": torch.tensor([10, 20, 30])},
... }
>>> # Slice the first 2 samples from all tensors
>>> bn.slice_along_batch(batch, stop=2)
{'inputs': tensor([[1, 2], [3, 4]]),
 'labels': tensor([0, 1]),
 'metadata': {'ids': tensor([10, 20])}}

Working with nested lists:

>>> data = [
...     torch.tensor([[1, 2], [3, 4]]),
...     torch.tensor([[5, 6], [7, 8]]),
... ]
>>> # Apply absolute value to all tensors
>>> bn.abs(data)
[tensor([[1, 2], [3, 4]]), tensor([[5, 6], [7, 8]])]

See Also

batchtensor.tensor: Similar functions for individual tensors.

batchtensor.nested.abs

abs(data: NestedTensor) -> NestedTensor

Return new tensors with the absolute value of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The absolute value of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import abs
>>> data = {
...     "a": torch.tensor([[-4, -3], [-2, -1], [0, 1], [2, 3], [4, 5]]),
...     "b": torch.tensor([2, 1, 0, -1, -2]),
... }
>>> out = abs(data)
>>> out
{'a': tensor([[4, 3], [2, 1], [0, 1], [2, 3], [4, 5]]), 'b': tensor([2, 1, 0, 1, 2])}

batchtensor.nested.acos

acos(data: NestedTensor) -> NestedTensor

Return new tensors with the inverse cosine of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The inverse cosine of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import acos
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = acos(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.acosh

acosh(data: NestedTensor) -> NestedTensor

Return new tensors with the inverse hyperbolic cosine of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The inverse hyperbolic cosine of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import acosh
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = acosh(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.amax_along_batch

amax_along_batch(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the maximum of all elements along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The maximum of all elements along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import amax_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = amax_along_batch(data)
>>> out
{'a': tensor([8, 9]), 'b': tensor(4)}
>>> out = amax_along_batch(data, keepdim=True)
>>> out
{'a': tensor([[8, 9]]), 'b': tensor([4])}

batchtensor.nested.amax_along_seq

amax_along_seq(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the maximum of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The maximum of all elements along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import amax_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = amax_along_seq(data)
>>> out
{'a': tensor([4, 9]), 'b': tensor([4])}
>>> out = amax_along_seq(data, keepdim=True)
>>> out
{'a': tensor([[4], [9]]), 'b': tensor([[4]])}

batchtensor.nested.amin_along_batch

amin_along_batch(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the minimum of all elements along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The minimum of all elements along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import amin_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = amin_along_batch(data)
>>> out
{'a': tensor([0, 1]), 'b': tensor(0)}
>>> out = amin_along_batch(data, keepdim=True)
>>> out
{'a': tensor([[0, 1]]), 'b': tensor([0])}

batchtensor.nested.amin_along_seq

amin_along_seq(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the minimum of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The minimum of all elements along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import amin_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = amin_along_seq(data)
>>> out
{'a': tensor([0, 5]), 'b': tensor([0])}
>>> out = amin_along_seq(data, keepdim=True)
>>> out
{'a': tensor([[0], [5]]), 'b': tensor([[0]])}

batchtensor.nested.argmax_along_batch

argmax_along_batch(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the indices of the maximum value of all elements along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The indices of the maximum value of all elements along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import argmax_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = argmax_along_batch(data)
>>> out
{'a': tensor([4, 4]), 'b': tensor(0)}
>>> out = argmax_along_batch(data, keepdim=True)
>>> out
{'a': tensor([[4, 4]]), 'b': tensor([0])}

batchtensor.nested.argmax_along_seq

argmax_along_seq(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the indices of the maximum value of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The indices of the maximum value of all elements along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import argmax_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = argmax_along_seq(data)
>>> out
{'a': tensor([4, 4]), 'b': tensor([0])}
>>> out = argmax_along_seq(data, keepdim=True)
>>> out
{'a': tensor([[4], [4]]), 'b': tensor([[0]])}

batchtensor.nested.argmin_along_batch

argmin_along_batch(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the indices of the minimum value of all elements along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The indices of the minimum value of all elements along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import argmin_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = argmin_along_batch(data)
>>> out
{'a': tensor([0, 0]), 'b': tensor(4)}
>>> out = argmin_along_batch(data, keepdim=True)
>>> out
{'a': tensor([[0, 0]]), 'b': tensor([4])}

batchtensor.nested.argmin_along_seq

argmin_along_seq(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the indices of the minimum value of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The indices of the minimum value of all elements along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import argmin_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = argmin_along_seq(data)
>>> out
{'a': tensor([0, 0]), 'b': tensor([4])}
>>> out = argmin_along_seq(data, keepdim=True)
>>> out
{'a': tensor([[0], [0]]), 'b': tensor([[4]])}

batchtensor.nested.argsort_along_batch

argsort_along_batch(
    data: NestedTensor,
    descending: bool = False,
    **kwargs: Any
) -> NestedTensor

Return the indices that sort a tensor along the batch dimension in ascending order by value.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
descending	bool	Controls the sorting order (ascending or descending).	False
kwargs	Any	Additional keywords arguments for torch.argsort.	{}

Returns:

Type	Description
NestedTensor	The indices that sort each tensor along the batch dimension

Example

>>> import torch
>>> from batchtensor.nested import argsort_along_batch
>>> data = {
...     "a": torch.tensor([[2, 6], [0, 3], [4, 9], [8, 1], [5, 7]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = argsort_along_batch(data)
>>> out
{'a': tensor([[1, 3], [0, 1], [2, 0], [4, 4], [3, 2]]), 'b': tensor([4, 3, 2, 1, 0])}
>>> out = argsort_along_batch(data, descending=True)
>>> out
{'a': tensor([[3, 2], [4, 4], [2, 0], [0, 1], [1, 3]]), 'b': tensor([0, 1, 2, 3, 4])}

batchtensor.nested.argsort_along_seq

argsort_along_seq(
    data: NestedTensor,
    descending: bool = False,
    **kwargs: Any
) -> NestedTensor

Return the indices that sort each tensor along the sequence dimension in ascending order by value.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
descending	bool	Controls the sorting order (ascending or descending).	False
kwargs	Any	Additional keywords arguments for torch.argsort.	{}

Returns:

Type	Description
NestedTensor	The indices that sort each tensor along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import argsort_along_seq
>>> data = {
...     "a": torch.tensor([[7, 3, 0, 8, 5], [1, 9, 6, 4, 2]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = argsort_along_seq(data)
>>> out
{'a': tensor([[2, 1, 4, 0, 3], [0, 4, 3, 2, 1]]), 'b': tensor([[4, 3, 2, 1, 0]])}
>>> out = argsort_along_seq(data, descending=True)
>>> out
{'a': tensor([[3, 0, 4, 1, 2], [1, 2, 3, 4, 0]]), 'b': tensor([[0, 1, 2, 3, 4]])}

batchtensor.nested.as_tensor

as_tensor(
    data: NestedTensor,
    dtype: dtype | None = None,
    device: device | None = None,
) -> NestedTensor

Create a new nested data structure with torch.Tensors.

This function recursively converts all array-like data (lists, tuples, numpy arrays, scalars, or existing tensors) to PyTorch tensors within a nested structure. Unlike torch.tensor(), this function shares memory with the input data when possible.

Note

This function preserves the structure of the input data while converting all array-like values to tensors. The output structure mirrors the input structure.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a value compatible with torch.as_tensor, such as lists, tuples, numpy arrays, Python scalars, or existing tensors. Can be a nested structure of dictionaries, lists, tuples containing such values.	required
dtype	dtype | None	The desired data type of returned tensors. If None, it infers data type from the input data. Common dtypes include torch.float32, torch.int64, torch.bool, etc.	None
device	device | None	The device of the constructed tensors (e.g., torch.device('cuda') or 'cpu'). If None and data contains a tensor, the device of that tensor is used. If None and data is not a tensor, the result tensor is constructed on the CPU.	None

Returns:

Type	Description
NestedTensor	A nested data structure with torch.Tensors. The output data has the same structure as the input, with all array-like values converted to tensors.

Example

>>> import numpy as np
>>> from batchtensor.nested import as_tensor
>>> # Convert mixed types in nested structure
>>> data = {"a": np.ones((2, 5), dtype=np.float32), "b": np.arange(5), "c": 42}
>>> out = as_tensor(data)
>>> out
{'a': tensor([[1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.]]),
 'b': tensor([0, 1, 2, 3, 4]),
 'c': tensor(42)}
>>> # Specify dtype for all tensors
>>> out = as_tensor({"values": np.array([1, 2, 3])}, dtype=torch.float32)
>>> out
{'values': tensor([1., 2., 3.])}

See Also

batchtensor.nested.from_numpy: Convert numpy arrays to tensors (shares memory). batchtensor.nested.to: Convert tensor dtype/device for existing tensors.

batchtensor.nested.asin

asin(data: NestedTensor) -> NestedTensor

Return new tensors with the arcsine of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The arcsine of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import asin
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = asin(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.asinh

asinh(data: NestedTensor) -> NestedTensor

Return new tensors with the inverse hyperbolic sine of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The inverse hyperbolic sine of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import asinh
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = asinh(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.atan

atan(data: NestedTensor) -> NestedTensor

Return new tensors with the arctangent of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The arctangent of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import atan
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = atan(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.atanh

atanh(data: NestedTensor) -> NestedTensor

Return new tensors with the inverse hyperbolic tangent of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The inverse hyperbolic tangent of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import atanh
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = atanh(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.cat_along_batch

cat_along_batch(
    data: Sequence[dict[Hashable, Tensor]],
) -> dict[Hashable, Tensor]

Concatenate the given tensors in the batch dimension.

All tensors must either have the same data type and shape (except in the concatenating dimension) or be empty.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	Sequence[dict[Hashable, Tensor]]	The input data to concatenate. The dictionaries must have the same keys.	required

Raises:

Type	Description
KeyError	if the dictionaries do not all have the same keys.

Returns:

Type	Description
dict[Hashable, Tensor]	The concatenated tensors along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import cat_along_batch
>>> data = [
...     {
...         "a": torch.tensor([[0, 1, 2], [4, 5, 6]]),
...         "b": torch.tensor([[10, 11, 12], [13, 14, 15]]),
...     },
...     {"a": torch.tensor([[7, 8, 9]]), "b": torch.tensor([[17, 18, 19]])},
... ]
>>> out = cat_along_batch(data)
>>> out
{'a': tensor([[0, 1, 2], [4, 5, 6], [7, 8, 9]]),
 'b': tensor([[10, 11, 12], [13, 14, 15], [17, 18, 19]])}

batchtensor.nested.cat_along_seq

cat_along_seq(
    data: Sequence[dict[Hashable, Tensor]],
) -> dict[Hashable, Tensor]

Concatenate the given tensors in the sequence dimension.

All tensors must either have the same data type and shape (except in the concatenating dimension) or be empty.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Sequence[dict[Hashable, Tensor]]	The input data to concatenate. The dictionaries must have the same keys.	required

Raises:

Type	Description
KeyError	if the dictionaries do not all have the same keys.

Returns:

Type	Description
dict[Hashable, Tensor]	The concatenated tensors along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import cat_along_seq
>>> data = [
...     {
...         "a": torch.tensor([[0, 1, 2], [4, 5, 6]]),
...         "b": torch.tensor([[10, 11, 12], [13, 14, 15]]),
...     },
...     {"a": torch.tensor([[7], [8]]), "b": torch.tensor([[17], [18]])},
... ]
>>> out = cat_along_seq(data)
>>> out
{'a': tensor([[0, 1, 2, 7], [4, 5, 6, 8]]),
 'b': tensor([[10, 11, 12, 17], [13, 14, 15, 18]])}

batchtensor.nested.chunk_along_batch

chunk_along_batch(
    data: dict[Hashable, Tensor], chunks: int
) -> tuple[dict[Hashable, Tensor], ...]

Split all the tensors into chunks along the batch dimension.

Each chunk is a view of the input tensor.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	dict[Hashable, Tensor]	The input data. Each item must be a tensor.	required
chunks	int	Number of chunks to return.	required

Returns:

Type	Description
tuple[dict[Hashable, Tensor], ...]	The data chunks.

Example

>>> import torch
>>> from batchtensor.nested import chunk_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> outputs = chunk_along_batch(data, chunks=3)
>>> outputs
({'a': tensor([[0, 1], [2, 3]]), 'b': tensor([4, 3])},
 {'a': tensor([[4, 5], [6, 7]]), 'b': tensor([2, 1])},
 {'a': tensor([[8, 9]]), 'b': tensor([0])})

batchtensor.nested.chunk_along_seq

chunk_along_seq(
    data: dict[Hashable, Tensor], chunks: int
) -> tuple[dict[Hashable, Tensor], ...]

Split all the tensors into chunks along the sequence dimension.

Each chunk is a view of the input tensor.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	dict[Hashable, Tensor]	The input data. Each item must be a tensor.	required
chunks	int	Number of chunks to return.	required

Returns:

Type	Description
tuple[dict[Hashable, Tensor], ...]	The data chunks.

Example

>>> import torch
>>> from batchtensor.nested import chunk_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> outputs = chunk_along_seq(data, chunks=3)
>>> outputs
({'a': tensor([[0, 1], [5, 6]]), 'b': tensor([[4, 3]])},
 {'a': tensor([[2, 3], [7, 8]]), 'b': tensor([[2, 1]])},
 {'a': tensor([[4], [9]]), 'b': tensor([[0]])})

batchtensor.nested.clamp

clamp(
    data: NestedTensor,
    min: float | None = None,
    max: float | None = None,
) -> NestedTensor

Clamp all elements in input into the range [min, max].

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
min	float | None	The lower-bound of the range to be clamped to.	None
max	float | None	The upper-bound of the range to be clamped to.	None

Returns:

Type	Description
NestedTensor	The clamp value of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import clamp
>>> data = {
...     "a": torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": torch.tensor([5, 4, 3, 2, 1]),
... }
>>> out = clamp(data, min=1, max=5)
>>> out
{'a': tensor([[1, 2], [3, 4], [5, 5], [5, 5], [5, 5]]), 'b': tensor([5, 4, 3, 2, 1])}

batchtensor.nested.clone

clone(data: NestedTensor) -> NestedTensor

Return a nested data structure with a copy of each tensor.

This function recursively applies torch.Tensor.clone() to all tensors in the nested data structure. Unlike the other functions in this module, the returned tensors do not share memory with the input tensors.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input nested data structure. Can be a dictionary, list, tuple, or any combination of these containing tensors. All leaf values in the structure must be tensors.	required

Returns:

Type	Description
NestedTensor	The data with each tensor replaced by a clone. The structure is preserved.

Example

>>> import torch
>>> from batchtensor.nested import clone
>>> data = {"a": torch.tensor([[0, 1], [2, 3]]), "b": torch.tensor([4, 5])}
>>> out = clone(data)
>>> out
{'a': tensor([[0, 1], [2, 3]]), 'b': tensor([4, 5])}

See Also

batchtensor.nested.detach: Detach tensors from the computation graph.

batchtensor.nested.contiguous

contiguous(
    data: NestedTensor, *args: Any, **kwargs: Any
) -> NestedTensor

Return a nested data structure with a contiguous copy of each tensor.

This function recursively applies torch.Tensor.contiguous() to all tensors in the nested data structure.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input nested data structure. Can be a dictionary, list, tuple, or any combination of these containing tensors. All leaf values in the structure must be tensors.	required
args	Any	Positional arguments passed to torch.Tensor.contiguous.	()
kwargs	Any	Keyword arguments passed to torch.Tensor.contiguous. Supports the memory_format argument.	{}

Returns:

Type	Description
NestedTensor	The data with each tensor made contiguous. The structure is preserved.

Example

>>> import torch
>>> from batchtensor.nested import contiguous
>>> data = {"a": torch.tensor([[0, 1], [2, 3]]), "b": torch.tensor([4, 5])}
>>> out = contiguous(data)
>>> out
{'a': tensor([[0, 1], [2, 3]]), 'b': tensor([4, 5])}

batchtensor.nested.cos

cos(data: NestedTensor) -> NestedTensor

Return new tensors with the cosine of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The cosine of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import cos
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = cos(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.cosh

cosh(data: NestedTensor) -> NestedTensor

Return new tensors with the hyperbolic cosine of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The hyperbolic cosine of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import cosh
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = cosh(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.cumprod_along_batch

cumprod_along_batch(data: NestedTensor) -> NestedTensor

Return the cumulative product of elements of input in the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The cumulative product of elements of input in the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import cumprod_along_batch
>>> data = {
...     "a": torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = cumprod_along_batch(data)
>>> out
{'a': tensor([[   1,    2], [   3,    8], [  15,   48], [ 105,  384], [ 945, 3840]]), 'b': tensor([ 4, 12, 24, 24,  0])}

batchtensor.nested.cumprod_along_seq

cumprod_along_seq(data: NestedTensor) -> NestedTensor

Return the cumulative product of elements of input in the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The cumulative product of elements of input in the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import cumprod_along_seq
>>> data = {
...     "a": torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = cumprod_along_seq(data)
>>> out
{'a': tensor([[    1,     2,     6,    24,   120], [    6,    42,   336,  3024, 30240]]), 'b': tensor([[ 4, 12, 24, 24,  0]])}

batchtensor.nested.cumsum_along_batch

cumsum_along_batch(data: NestedTensor) -> NestedTensor

Return the cumulative sum of elements of input in the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The cumulative sum of elements of input in the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import cumsum_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = cumsum_along_batch(data)
>>> out
{'a': tensor([[ 0,  1], [ 2,  4], [ 6,  9], [12, 16], [20, 25]]), 'b': tensor([ 4,  7,  9, 10, 10])}

batchtensor.nested.cumsum_along_seq

cumsum_along_seq(data: NestedTensor) -> NestedTensor

Return the cumulative sum of elements of input in the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The cumulative sum of elements of input in the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import cumsum_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = cumsum_along_seq(data)
>>> out
{'a': tensor([[ 0,  1,  3,  6, 10], [ 5, 11, 18, 26, 35]]), 'b': tensor([[ 4,  7,  9, 10, 10]])}

batchtensor.nested.detach

detach(data: NestedTensor) -> NestedTensor

Return a nested data structure with each tensor detached from the current computation graph.

This function recursively applies torch.Tensor.detach() to all tensors in the nested data structure. The returned tensors share memory with the input tensors.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input nested data structure. Can be a dictionary, list, tuple, or any combination of these containing tensors. All leaf values in the structure must be tensors.	required

Returns:

Type	Description
NestedTensor	The data with each tensor detached. The structure is preserved.

Example

>>> import torch
>>> from batchtensor.nested import detach
>>> data = {
...     "a": torch.tensor([[0.0, 1.0], [2.0, 3.0]], requires_grad=True),
...     "b": torch.tensor([4.0, 5.0]),
... }
>>> out = detach(data)
>>> out["a"].requires_grad
False

See Also

batchtensor.nested.clone: Create an independent copy of each tensor.

batchtensor.nested.exp

exp(data: NestedTensor) -> NestedTensor

Return new tensors with the exponential of the elements.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The exponential of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import exp
>>> data = {
...     "a": torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": torch.tensor([5, 4, 3, 2, 1]),
... }
>>> out = exp(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.exp2

exp2(data: NestedTensor) -> NestedTensor

Return new tensors with the base two exponential of the elements.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The base two exponential of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import exp2
>>> data = {
...     "a": torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": torch.tensor([5, 4, 3, 2, 1]),
... }
>>> out = exp2(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.expm1

expm1(data: NestedTensor) -> NestedTensor

Return new tensors with the exponential of the elements minus 1.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The exponential of the elements minus 1. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import expm1
>>> data = {
...     "a": torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": torch.tensor([5, 4, 3, 2, 1]),
... }
>>> out = expm1(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.from_numpy

from_numpy(data: NestedTensor) -> NestedTensor

Create a new nested data structure where the numpy.ndarrays are converted to torch.Tensors.

This function recursively converts all numpy arrays in a nested structure to PyTorch tensors. The conversion uses torch.from_numpy(), which creates tensors that share memory with the original arrays.

Note

The returned torch.Tensors and numpy.ndarrays share the same underlying memory. Modifications to the torch.Tensors will be reflected in the numpy.ndarrays and vice versa. This is efficient but requires caution when modifying data.

Warning

Since memory is shared, modifying the tensor will modify the original numpy array. If you need independent copies, use torch.tensor() or call .clone() on the result.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item should be a numpy.ndarray. Can be a nested structure of dictionaries, lists, tuples containing numpy arrays.	required

Returns:

Type	Description
NestedTensor	A nested data structure with torch.Tensors instead of numpy.ndarrays. The output data has the same structure as the input, and the tensors share memory with the original arrays.

Example

>>> import numpy as np
>>> from batchtensor.nested import from_numpy
>>> data = {"a": np.ones((2, 5), dtype=np.float32), "b": np.arange(5)}
>>> out = from_numpy(data)
>>> out
{'a': tensor([[1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.]]), 'b': tensor([0, 1, 2, 3, 4])}
>>> # Demonstrate memory sharing
>>> data["a"][0, 0] = 999
>>> out["a"][0, 0]  # Reflects the change
tensor(999.)

See Also

batchtensor.nested.to_numpy: Convert tensors to numpy arrays (shares memory). batchtensor.nested.as_tensor: More flexible conversion supporting various input types.

batchtensor.nested.index_select_along_batch

index_select_along_batch(
    data: NestedTensor, index: Tensor
) -> NestedTensor

Return the tensors which indexes the input tensor along the batch dimension using the entries in index which is a LongTensor.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
index	Tensor	The 1-D tensor containing the indices to index.	required

Returns:

Type	Description
NestedTensor	The indexed tensors along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import index_select_along_batch
>>> tensors = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = index_select_along_batch(tensors, torch.tensor([2, 4]))
>>> out
{'a': tensor([[4, 5], [8, 9]]), 'b': tensor([2, 0])}
>>> out = index_select_along_batch(tensors, torch.tensor([4, 3, 2, 1, 0]))
>>> out
{'a': tensor([[8, 9], [6, 7], [4, 5], [2, 3], [0, 1]]), 'b': tensor([0, 1, 2, 3, 4])}

batchtensor.nested.index_select_along_seq

index_select_along_seq(
    data: NestedTensor, index: Tensor
) -> NestedTensor

Return the tensors which indexes the input tensor along the sequence dimension using the entries in index which is a LongTensor.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
index	Tensor	The 1-D tensor containing the indices to index.	required

Returns:

Type	Description
NestedTensor	The indexed tensors along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import index_select_along_seq
>>> tensors = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = index_select_along_seq(tensors, torch.tensor([2, 4]))
>>> out
{'a': tensor([[2, 4], [7, 9]]), 'b': tensor([[2, 0]])}
>>> out = index_select_along_seq(tensors, torch.tensor([4, 3, 2, 1, 0]))
>>> out
{'a': tensor([[4, 3, 2, 1, 0], [9, 8, 7, 6, 5]]), 'b': tensor([[0, 1, 2, 3, 4]])}

batchtensor.nested.log

log(data: NestedTensor) -> NestedTensor

Return new tensors with the natural logarithm of the elements.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The natural logarithm of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import log
>>> data = {
...     "a": torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": torch.tensor([5, 4, 3, 2, 1]),
... }
>>> out = log(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.log10

log10(data: NestedTensor) -> NestedTensor

Return new tensors with the logarithm to the base 10 of the elements.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The logarithm to the base 10 of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import log10
>>> data = {
...     "a": torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": torch.tensor([5, 4, 3, 2, 1]),
... }
>>> out = log10(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.log1p

log1p(data: NestedTensor) -> NestedTensor

Return new tensors with the natural logarithm of (1 + input).

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The natural logarithm of (1 + input). The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import log1p
>>> data = {
...     "a": torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": torch.tensor([5, 4, 3, 2, 1]),
... }
>>> out = log1p(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.log2

log2(data: NestedTensor) -> NestedTensor

Return new tensors with the logarithm to the base 2 of the elements.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The logarithm to the base 2 of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import log2
>>> data = {
...     "a": torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": torch.tensor([5, 4, 3, 2, 1]),
... }
>>> out = log2(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.max_along_batch

max_along_batch(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the maximum of all elements along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The first tensor will be populated with the maximum values and the second tensor, which must have dtype long, with their indices in the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import max_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = max_along_batch(data)
>>> out
{'a': torch.return_types.max(
values=tensor([8, 9]),
indices=tensor([4, 4])),
'b': torch.return_types.max(
values=tensor(4),
indices=tensor(0))}
>>> out = max_along_batch(data, keepdim=True)
>>> out
{'a': torch.return_types.max(
values=tensor([[8, 9]]),
indices=tensor([[4, 4]])),
'b': torch.return_types.max(
values=tensor([4]),
indices=tensor([0]))}

batchtensor.nested.max_along_seq

max_along_seq(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the maximum of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The first tensor will be populated with the maximum values and the second tensor, which must have dtype long, with their indices in the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import max_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = max_along_seq(data)
>>> out
{'a': torch.return_types.max(
values=tensor([4, 9]),
indices=tensor([4, 4])),
'b': torch.return_types.max(
values=tensor([4]),
indices=tensor([0]))}
>>> out = max_along_seq(data, keepdim=True)
>>> out
{'a': torch.return_types.max(
values=tensor([[4], [9]]),
indices=tensor([[4], [4]])),
'b': torch.return_types.max(
values=tensor([[4]]),
indices=tensor([[0]]))}

batchtensor.nested.mean_along_batch

mean_along_batch(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the mean of all elements along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The mean of all elements along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import mean_along_batch
>>> data = {
...     "a": torch.tensor([[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0], [8.0, 9.0]]),
...     "b": torch.tensor([4, 3, 2, 1, 0], dtype=torch.float),
... }
>>> out = mean_along_batch(data)
>>> out
{'a': tensor([4., 5.]), 'b': tensor(2.)}
>>> out = mean_along_batch(data, keepdim=True)
>>> out
{'a': tensor([[4., 5.]]), 'b': tensor([2.])}

batchtensor.nested.mean_along_seq

mean_along_seq(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the mean of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The mean of all elements along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import mean_along_seq
>>> data = {
...     "a": torch.tensor([[0.0, 1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0, 9.0]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]], dtype=torch.float),
... }
>>> out = mean_along_seq(data)
>>> out
{'a': tensor([2., 7.]), 'b': tensor([2.])}
>>> out = mean_along_seq(data, keepdim=True)
>>> out
{'a': tensor([[2.], [7.]]), 'b': tensor([[2.]])}

batchtensor.nested.median_along_batch

median_along_batch(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the median of all elements along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The first tensor will be populated with the median values and the second tensor, which must have dtype long, with their indices in the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import median_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = median_along_batch(data)
>>> out
{'a': torch.return_types.median(
values=tensor([4, 5]),
indices=tensor([2, 2])),
'b': torch.return_types.median(
values=tensor(2),
indices=tensor(2))}
>>> out = median_along_batch(data, keepdim=True)
>>> out
{'a': torch.return_types.median(
values=tensor([[4, 5]]),
indices=tensor([[2, 2]])),
'b': torch.return_types.median(
values=tensor([2]),
indices=tensor([2]))}

batchtensor.nested.median_along_seq

median_along_seq(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the median of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The first tensor will be populated with the median values and the second tensor, which must have dtype long, with their indices in the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import median_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = median_along_seq(data)
>>> out
{'a': torch.return_types.median(
values=tensor([2, 7]),
indices=tensor([2, 2])),
'b': torch.return_types.median(
values=tensor([2]),
indices=tensor([2]))}
>>> out = median_along_seq(data, keepdim=True)
>>> out
{'a': torch.return_types.median(
values=tensor([[2], [7]]),
indices=tensor([[2], [2]])),
'b': torch.return_types.median(
values=tensor([[2]]),
indices=tensor([[2]]))}

batchtensor.nested.min_along_batch

min_along_batch(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the minimum of all elements along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The first tensor will be populated with the minimum values and the second tensor, which must have dtype long, with their indices in the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import min_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = min_along_batch(data)
>>> out
{'a': torch.return_types.min(
values=tensor([0, 1]),
indices=tensor([0, 0])),
'b': torch.return_types.min(
values=tensor(0),
indices=tensor(4))}
>>> out = min_along_batch(data, keepdim=True)
>>> out
{'a': torch.return_types.min(
values=tensor([[0, 1]]),
indices=tensor([[0, 0]])),
'b': torch.return_types.min(
values=tensor([0]),
indices=tensor([4]))}

batchtensor.nested.min_along_seq

min_along_seq(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the minimum of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The first tensor will be populated with the minimum values and the second tensor, which must have dtype long, with their indices in the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import min_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = min_along_seq(data)
>>> out
{'a': torch.return_types.min(
values=tensor([0, 5]),
indices=tensor([0, 0])),
'b': torch.return_types.min(
values=tensor([0]),
indices=tensor([4]))}
>>> out = min_along_seq(data, keepdim=True)
>>> out
{'a': torch.return_types.min(
values=tensor([[0], [5]]),
indices=tensor([[0], [0]])),
'b': torch.return_types.min(
values=tensor([[0]]),
indices=tensor([[4]]))}

batchtensor.nested.pad_along_batch

pad_along_batch(
    data: NestedTensor, pad_size: int, value: float = 0.0
) -> NestedTensor

Pad all the tensors along the batch dimension.

The padding is added at the end of the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
pad_size	int	The number of items to add to the batch dimension. Must be a non-negative integer.	required
value	float	The fill value used for the padded items.	0.0

Returns:

Type	Description
NestedTensor	The padded tensors.

Example

>>> import torch
>>> from batchtensor.nested import pad_along_batch
>>> data = {"a": torch.tensor([[0, 1], [2, 3]]), "b": torch.tensor([4, 5])}
>>> out = pad_along_batch(data, pad_size=1)
>>> out
{'a': tensor([[0, 1], [2, 3], [0, 0]]), 'b': tensor([4, 5, 0])}

See Also

pad_along_seq: Pad along the sequence dimension instead.

batchtensor.nested.pad_along_seq

pad_along_seq(
    data: NestedTensor, pad_size: int, value: float = 0.0
) -> NestedTensor

Pad all the tensors along the sequence dimension.

The padding is added at the end of the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
pad_size	int	The number of items to add to the sequence dimension. Must be a non-negative integer.	required
value	float	The fill value used for the padded items.	0.0

Returns:

Type	Description
NestedTensor	The padded tensors.

Example

>>> import torch
>>> from batchtensor.nested import pad_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2], [3, 4, 5]]),
...     "b": torch.tensor([[6, 7, 8]]),
... }
>>> out = pad_along_seq(data, pad_size=1)
>>> out
{'a': tensor([[0, 1, 2, 0], [3, 4, 5, 0]]), 'b': tensor([[6, 7, 8, 0]])}

See Also

pad_along_batch: Pad along the batch dimension instead. lengths_to_mask: Build a mask that flags the padded positions.

batchtensor.nested.permute_along_batch

permute_along_batch(
    data: NestedTensor, permutation: Tensor
) -> NestedTensor

Permute all the tensors along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
permutation	Tensor	The 1-D tensor containing the indices of the permutation. The shape should match the batch dimension of the tensor.	required

Returns:

Type	Description
NestedTensor	The data with permuted tensors along the batch dimension. The output data has the same structure as the input data.

Raises:

Type	Description
RuntimeError	if the shape of the permutation does not match the batch dimension of the tensor.

Example

>>> import torch
>>> from batchtensor.nested import permute_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = permute_along_batch(data, torch.tensor([2, 1, 3, 0, 4]))
>>> out
{'a': tensor([[4, 5], [2, 3], [6, 7], [0, 1], [8, 9]]), 'b': tensor([2, 3, 1, 4, 0])}

batchtensor.nested.permute_along_seq

permute_along_seq(
    data: NestedTensor, permutation: Tensor
) -> NestedTensor

Permute all the tensors along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
permutation	Tensor	The 1-D tensor containing the indices of the permutation. The shape should match the sequence dimension of the tensor.	required

Returns:

Type	Description
NestedTensor	The data with permuted tensors along the sequence dimension. The output data has the same structure as the input data.

Raises:

Type	Description
RuntimeError	if the shape of the permutation does not match the sequence dimension of the tensor.

Example

>>> import torch
>>> from batchtensor.nested import permute_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = permute_along_seq(data, torch.tensor([2, 1, 3, 0, 4]))
>>> out
{'a': tensor([[2, 1, 3, 0, 4], [7, 6, 8, 5, 9]]), 'b': tensor([[2, 3, 1, 4, 0]])}

batchtensor.nested.pin_memory

pin_memory(data: NestedTensor) -> NestedTensor

Return a nested data structure with each tensor copied to pinned memory.

This function recursively applies torch.Tensor.pin_memory() to all tensors in the nested data structure. Pinned (page-locked) CPU tensors allow faster host-to-device transfers.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input nested data structure. Can be a dictionary, list, tuple, or any combination of these containing tensors. All leaf values in the structure must be CPU tensors.	required

Returns:

Type	Description
NestedTensor	The data with each tensor copied to pinned memory. The structure is preserved.

Example

>>> import torch
>>> from batchtensor.nested import pin_memory
>>> data = {"a": torch.tensor([[0, 1], [2, 3]]), "b": torch.tensor([4, 5])}
>>> out = pin_memory(data)  # doctest: +SKIP

batchtensor.nested.prod_along_batch

prod_along_batch(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the product of all elements along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The product of all elements along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import prod_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([5, 4, 3, 2, 1]),
... }
>>> out = prod_along_batch(data)
>>> out
{'a': tensor([  0, 945]), 'b': tensor(120)}
>>> out = prod_along_batch(data, keepdim=True)
>>> out
{'a': tensor([[  0, 945]]), 'b': tensor([120])}

batchtensor.nested.prod_along_seq

prod_along_seq(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the product of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The product of all elements along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import prod_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[5, 4, 3, 2, 1]]),
... }
>>> out = prod_along_seq(data)
>>> out
{'a': tensor([    0, 15120]), 'b': tensor([120])}
>>> out = prod_along_seq(data, keepdim=True)
>>> out
{'a': tensor([[    0], [15120]]), 'b': tensor([[120]])}

batchtensor.nested.repeat_along_seq

repeat_along_seq(
    data: NestedTensor, repeats: int
) -> NestedTensor

Repeat all the tensors along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
repeats	int	The number of times to repeat the data along the sequence dimension.	required

Returns:

Type	Description
NestedTensor	The tensors repeated along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import repeat_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = repeat_along_seq(data, 2)
>>> out
{'a': tensor([[0, 1, 2, 3, 4, 0, 1, 2, 3, 4], [5, 6, 7, 8, 9, 5, 6, 7, 8, 9]]),
 'b': tensor([[4, 3, 2, 1, 0, 4, 3, 2, 1, 0]])}

batchtensor.nested.select_along_batch

select_along_batch(
    data: NestedTensor, index: int
) -> NestedTensor

Slice the tensors along the batch dimension at the given index.

This function returns a view of the original tensor with the batch dimension removed.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
index	int	The index to select with.	required

Returns:

Type	Description
NestedTensor	The sliced tensors along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import select_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = select_along_batch(data, index=2)
>>> out
{'a': tensor([4, 5]), 'b': tensor(2)}

batchtensor.nested.select_along_seq

select_along_seq(
    data: NestedTensor, index: int
) -> NestedTensor

Slice the tensors along the sequence dimension at the given index.

This function returns a view of the original tensor with the sequence dimension removed.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
index	int	The index to select with.	required

Returns:

Type	Description
NestedTensor	The sliced tensors along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import select_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = select_along_seq(data, index=2)
>>> out
{'a': tensor([2, 7]), 'b': tensor([2])}

batchtensor.nested.shuffle_along_batch

shuffle_along_batch(
    data: NestedTensor, generator: Generator | None = None
) -> NestedTensor

Shuffle all the tensors along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
generator	Generator | None	An optional random number generator.	None

Returns:

Type	Description
NestedTensor	The data with shuffled tensors along the batch dimension. The output data has the same structure as the input data.

Example

>>> import torch
>>> from batchtensor.nested import shuffle_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = shuffle_along_batch(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.shuffle_along_seq

shuffle_along_seq(
    data: NestedTensor, generator: Generator | None = None
) -> NestedTensor

Shuffle all the tensors along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
generator	Generator | None	An optional random number generator.	None

Returns:

Type	Description
NestedTensor	The data with shuffled tensors along the sequence dimension. The output data has the same structure as the input data.

Example

>>> import torch
>>> from batchtensor.nested import shuffle_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = shuffle_along_seq(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([[...]])}

batchtensor.nested.sin

sin(data: NestedTensor) -> NestedTensor

Return new tensors with the sine of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The sine of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import sin
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = sin(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.sinh

sinh(data: NestedTensor) -> NestedTensor

Return new tensors with the hyperbolic sine of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The hyperbolic sine of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import sinh
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = sinh(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.slice_along_batch

slice_along_batch(
    data: NestedTensor,
    start: int = 0,
    stop: int | None = None,
    step: int = 1,
) -> NestedTensor

Slice all the tensors along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
start	int	The index where the slicing of object starts.	0
stop	int | None	The index where the slicing of object stops. None means last.	None
step	int	The increment between each index for slicing.	1

Returns:

Type	Description
NestedTensor	The sliced tensor along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import slice_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = slice_along_batch(data, start=2)
>>> out
{'a': tensor([[4, 5], [6, 7], [8, 9]]), 'b': tensor([2, 1, 0])}
>>> out = slice_along_batch(data, stop=3)
>>> out
{'a': tensor([[0, 1], [2, 3], [4, 5]]), 'b': tensor([4, 3, 2])}
>>> out = slice_along_batch(data, step=2)
>>> out
{'a': tensor([[0, 1], [4, 5], [8, 9]]), 'b': tensor([4, 2, 0])}

batchtensor.nested.slice_along_seq

slice_along_seq(
    data: NestedTensor,
    start: int = 0,
    stop: int | None = None,
    step: int = 1,
) -> NestedTensor

Slice all the tensors along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
start	int	The index where the slicing of object starts.	0
stop	int | None	The index where the slicing of object stops. None means last.	None
step	int	The increment between each index for slicing.	1

Returns:

Type	Description
NestedTensor	The sliced tensor along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import slice_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = slice_along_seq(data, start=2)
>>> out
{'a': tensor([[2, 3, 4], [7, 8, 9]]), 'b': tensor([[2, 1, 0]])}
>>> out = slice_along_seq(data, stop=3)
>>> out
{'a': tensor([[0, 1, 2], [5, 6, 7]]), 'b': tensor([[4, 3, 2]])}
>>> out = slice_along_seq(data, step=2)
>>> out
{'a': tensor([[0, 2, 4], [5, 7, 9]]), 'b': tensor([[4, 2, 0]])}

batchtensor.nested.sort_along_batch

sort_along_batch(
    data: NestedTensor,
    descending: bool = False,
    **kwargs: Any
) -> NestedTensor

Sort the elements of the input tensor along the batch dimension in ascending order by value.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
descending	bool	Controls the sorting order (ascending or descending).	False
kwargs	Any	Additional keywords arguments for torch.sort.	{}

Returns:

Type	Description
NestedTensor	A similar object where each tensor is replaced by a namedtuple of (values, indices), where the values are the sorted values and indices are the indices of the elements in the original input tensor.

Example

>>> import torch
>>> from batchtensor.nested import sort_along_batch
>>> data = {
...     "a": torch.tensor([[2, 6], [0, 3], [4, 9], [8, 1], [5, 7]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = sort_along_batch(data)
>>> out
{'a': torch.return_types.sort(
values=tensor([[0, 1], [2, 3], [4, 6], [5, 7], [8, 9]]),
indices=tensor([[1, 3], [0, 1], [2, 0], [4, 4], [3, 2]])),
'b': torch.return_types.sort(
values=tensor([0, 1, 2, 3, 4]),
indices=tensor([4, 3, 2, 1, 0]))}
>>> out = sort_along_batch(data, descending=True)
>>> out
{'a': torch.return_types.sort(
values=tensor([[8, 9], [5, 7], [4, 6], [2, 3], [0, 1]]),
indices=tensor([[3, 2], [4, 4], [2, 0], [0, 1], [1, 3]])),
'b': torch.return_types.sort(
values=tensor([4, 3, 2, 1, 0]),
indices=tensor([0, 1, 2, 3, 4]))}

batchtensor.nested.sort_along_seq

sort_along_seq(
    data: NestedTensor,
    descending: bool = False,
    **kwargs: Any
) -> NestedTensor

Sort the elements of the input tensor along the sequence dimension in ascending order by value.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
descending	bool	Controls the sorting order (ascending or descending).	False
kwargs	Any	Additional keywords arguments for torch.sort.	{}

Returns:

Type	Description
NestedTensor	A similar object where each tensor is replaced by a namedtuple of (values, indices), where the values are the sorted values and indices are the indices of the elements in the original input tensor.

Example

>>> import torch
>>> from batchtensor.nested import sort_along_seq
>>> data = {
...     "a": torch.tensor([[7, 3, 0, 8, 5], [1, 9, 6, 4, 2]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = sort_along_seq(data)
>>> out
{'a': torch.return_types.sort(
values=tensor([[0, 3, 5, 7, 8], [1, 2, 4, 6, 9]]),
indices=tensor([[2, 1, 4, 0, 3], [0, 4, 3, 2, 1]])),
'b': torch.return_types.sort(
values=tensor([[0, 1, 2, 3, 4]]),
indices=tensor([[4, 3, 2, 1, 0]]))}
>>> out = sort_along_seq(data, descending=True)
>>> out
{'a': torch.return_types.sort(
values=tensor([[8, 7, 5, 3, 0], [9, 6, 4, 2, 1]]),
indices=tensor([[3, 0, 4, 1, 2], [1, 2, 3, 4, 0]])),
'b': torch.return_types.sort(
values=tensor([[4, 3, 2, 1, 0]]),
indices=tensor([[0, 1, 2, 3, 4]]))}

batchtensor.nested.split_along_batch

split_along_batch(
    data: dict[Hashable, Tensor],
    split_size_or_sections: int | Sequence[int],
) -> tuple[dict[Hashable, Tensor], ...]

Split all the tensors into chunks along the batch dimension.

Each chunk is a view of the original tensor.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	dict[Hashable, Tensor]	The input data. Each item must be a tensor.	required
split_size_or_sections	int | Sequence[int]	Size of a single chunk or list of sizes for each chunk	required

Returns:

Type	Description
tuple[dict[Hashable, Tensor], ...]	The data chunks.

Example

>>> import torch
>>> from batchtensor.nested import split_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> outputs = split_along_batch(data, split_size_or_sections=2)
>>> outputs
({'a': tensor([[0, 1], [2, 3]]), 'b': tensor([4, 3])},
 {'a': tensor([[4, 5], [6, 7]]), 'b': tensor([2, 1])},
 {'a': tensor([[8, 9]]), 'b': tensor([0])})

batchtensor.nested.split_along_seq

split_along_seq(
    data: dict[Hashable, Tensor],
    split_size_or_sections: int | Sequence[int],
) -> tuple[dict[Hashable, Tensor], ...]

Split all the tensors into chunks along the sequence dimension.

Each chunk is a view of the original tensor.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	dict[Hashable, Tensor]	The input data. Each item must be a tensor.	required
split_size_or_sections	int | Sequence[int]	Size of a single chunk or list of sizes for each chunk	required

Returns:

Type	Description
tuple[dict[Hashable, Tensor], ...]	The data chunks.

Example

>>> import torch
>>> from batchtensor.nested import split_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> outputs = split_along_seq(data, split_size_or_sections=2)
>>> outputs
({'a': tensor([[0, 1], [5, 6]]), 'b': tensor([[4, 3]])},
 {'a': tensor([[2, 3], [7, 8]]), 'b': tensor([[2, 1]])},
 {'a': tensor([[4], [9]]), 'b': tensor([[0]])})

batchtensor.nested.stack_along_batch

stack_along_batch(
    data: Sequence[dict[Hashable, Tensor]],
) -> dict[Hashable, Tensor]

Stack the given samples along a new batch dimension.

This is the nested equivalent of a collate function: given a sequence of samples (dictionaries of tensors without a batch dimension), it combines them into a single batched dictionary. Each tensor must have the same shape across all the samples.

Parameters:

Name	Type	Description	Default
data	Sequence[dict[Hashable, Tensor]]	The input samples to stack. The dictionaries must have the same keys.	required

Raises:

Type	Description
KeyError	if the dictionaries do not all have the same keys.

Returns:

Type	Description
dict[Hashable, Tensor]	The stacked tensors along a new batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import stack_along_batch
>>> data = [
...     {"a": torch.tensor([0, 1, 2]), "b": torch.tensor(1)},
...     {"a": torch.tensor([3, 4, 5]), "b": torch.tensor(2)},
... ]
>>> out = stack_along_batch(data)
>>> out
{'a': tensor([[0, 1, 2], [3, 4, 5]]), 'b': tensor([1, 2])}

See Also

cat_along_batch: Concatenate along an existing batch dimension. stack_along_seq: Stack along a new sequence dimension instead.

batchtensor.nested.stack_along_seq

stack_along_seq(
    data: Sequence[dict[Hashable, Tensor]],
) -> dict[Hashable, Tensor]

Stack the given samples along a new sequence dimension.

Combines a sequence of per-step samples (dictionaries of tensors without a sequence dimension) into a single dictionary with a sequence dimension. Each tensor must have the same shape across all the samples.

Parameters:

Name	Type	Description	Default
data	Sequence[dict[Hashable, Tensor]]	The input samples to stack. The dictionaries must have the same keys.	required

Raises:

Type	Description
KeyError	if the dictionaries do not all have the same keys.

Returns:

Type	Description
dict[Hashable, Tensor]	The stacked tensors along a new sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import stack_along_seq
>>> data = [
...     {"a": torch.tensor([0, 1]), "b": torch.tensor([10, 11])},
...     {"a": torch.tensor([2, 3]), "b": torch.tensor([12, 13])},
... ]
>>> out = stack_along_seq(data)
>>> out
{'a': tensor([[0, 2], [1, 3]]), 'b': tensor([[10, 12], [11, 13]])}

See Also

cat_along_seq: Concatenate along an existing sequence dimension. stack_along_batch: Stack along a new batch dimension instead.

batchtensor.nested.sum_along_batch

sum_along_batch(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the sum of all elements along the batch dimension.

Note

This function assumes the batch dimension is the first dimension of the tensors. All the tensors should have the same batch size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The sum of all elements along the batch dimension.

Example

>>> import torch
>>> from batchtensor.nested import sum_along_batch
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> out = sum_along_batch(data)
>>> out
{'a': tensor([20, 25]), 'b': tensor(10)}
>>> out = sum_along_batch(data, keepdim=True)
>>> out
{'a': tensor([[20, 25]]), 'b': tensor([10])}

batchtensor.nested.sum_along_seq

sum_along_seq(
    data: NestedTensor, keepdim: bool = False
) -> NestedTensor

Return the sum of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension of the tensors. All the tensors should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required
keepdim	bool	Whether the output tensor has dim retained or not.	False

Returns:

Type	Description
NestedTensor	The sum of all elements along the sequence dimension.

Example

>>> import torch
>>> from batchtensor.nested import sum_along_seq
>>> data = {
...     "a": torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": torch.tensor([[4, 3, 2, 1, 0]]),
... }
>>> out = sum_along_seq(data)
>>> out
{'a': tensor([10, 35]), 'b': tensor([10])}
>>> out = sum_along_seq(data, keepdim=True)
>>> out
{'a': tensor([[10], [35]]), 'b': tensor([[10]])}

batchtensor.nested.tan

tan(data: NestedTensor) -> NestedTensor

Return new tensors with the tangent of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The tangent of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import tan
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = tan(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.tanh

tanh(data: NestedTensor) -> NestedTensor

Return new tensors with the hyperbolic tangent of each element.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a tensor.	required

Returns:

Type	Description
NestedTensor	The hyperbolic tangent of the elements. The output has the same structure as the input.

Example

>>> import torch
>>> from batchtensor.nested import tanh
>>> data = {"a": torch.randn(5, 2), "b": torch.rand(5)}
>>> out = tanh(data)
>>> out
{'a': tensor([[...]]), 'b': tensor([...])}

batchtensor.nested.to

to(
    data: NestedTensor, *args: Any, **kwargs: Any
) -> NestedTensor

Perform Tensor dtype and/or device conversion on all tensors in nested data.

This function recursively applies torch.Tensor.to() to all tensors in the nested data structure, allowing you to convert dtypes, move to different devices, or change other tensor properties for all tensors at once.

Note

This function preserves the structure of the input data while converting all tensors within it.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input nested data structure. Can be a dictionary, list, tuple, or any combination of these containing tensors. All leaf values in the structure must be tensors.	required
args	Any	Positional arguments passed to torch.Tensor.to. Common usage includes passing a device (e.g., torch.device('cuda')), dtype (e.g., torch.float32), or another tensor to match device and dtype.	()
kwargs	Any	Keyword arguments passed to torch.Tensor.to. Supports arguments like dtype, device, non_blocking, copy, and memory_format.	{}

Returns:

Type	Description
NestedTensor	The data after conversion. The structure is preserved, with all tensors converted according to the specified arguments.

Example

>>> import torch
>>> from batchtensor.nested import to
>>> data = {
...     "a": torch.tensor([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": torch.tensor([4, 3, 2, 1, 0]),
... }
>>> # Convert to float dtype
>>> out = to(data, dtype=torch.float)
>>> out
{'a': tensor([[0., 1.], [2., 3.], [4., 5.], [6., 7.], [8., 9.]]),
 'b': tensor([4., 3., 2., 1., 0.])}
>>> # Move to GPU (if available) with float32 dtype
>>> # out = to(data, device='cuda', dtype=torch.float32)

See Also

batchtensor.nested.as_tensor: Convert data to tensor format. batchtensor.nested.from_numpy: Convert numpy arrays to tensors.

batchtensor.nested.to_numpy

to_numpy(data: NestedTensor) -> NestedTensor

Create a new nested data structure where the torch.Tensors are converted to numpy.ndarrays.

This function recursively converts all PyTorch tensors in a nested structure to numpy arrays. The conversion uses tensor.numpy(), which creates arrays that share memory with the original tensors when possible (for tensors on CPU with compatible dtypes).

Note

The returned torch.Tensors and numpy.ndarrays share the same underlying memory when tensors are on CPU with compatible dtypes. Modifications to the torch.Tensors will be reflected in the numpy.ndarrays and vice versa.

Warning

• Tensors on CUDA devices will be copied to CPU before conversion.
• Tensors with gradients enabled will raise an error unless you call .detach() first.
• Since memory is shared (for CPU tensors), modifying the array will modify the original tensor.

Parameters:

Name	Type	Description	Default
data	NestedTensor	The input data. Each item must be a torch.Tensor. Can be a nested structure of dictionaries, lists, tuples containing tensors. All tensors should be on CPU without gradients for efficient conversion.	required

Returns:

Type	Description
NestedTensor	A nested data structure with numpy.ndarrays instead of torch.Tensors. The output data has the same structure as the input. For CPU tensors with compatible dtypes, the arrays share memory with the original tensors.

Example

>>> import torch
>>> from batchtensor.nested import to_numpy
>>> data = {"a": torch.ones(2, 5), "b": torch.tensor([0, 1, 2, 3, 4])}
>>> out = to_numpy(data)
>>> out
{'a': array([[1., 1., 1., 1., 1.], [1., 1., 1., 1., 1.]], dtype=float32), 'b': array([0, 1, 2, 3, 4])}
>>> # Demonstrate memory sharing for CPU tensors
>>> data["a"][0, 0] = 999
>>> out["a"][0, 0]  # Reflects the change
np.float32(999.0)

See Also

batchtensor.nested.from_numpy: Convert numpy arrays to tensors (shares memory). batchtensor.nested.as_tensor: Convert various data types to tensors.