Nested

batchtensor.nested

Contain functions to manipulate nested data.

batchtensor.nested.abs

abs(data: Any) -> Any

Return new tensors with the absolute value of each element.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

Return new tensors with the inverse cosine of each element.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

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

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The maximum of all elements along the batch dimension.

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The maximum of all elements along the sequence dimension.

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The minimum of all elements along the batch dimension.

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The minimum of all elements along the sequence dimension.

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The indices of the maximum value of all elements along the batch dimension.

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The indices of the maximum value of all elements along the sequence dimension.

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The indices of the minimum value of all elements along the batch dimension.

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The indices of the minimum value of all elements along the sequence dimension.

Example usage:

>>> 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: Any, descending: bool = False, **kwargs: Any
) -> Any

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	Any	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
Any	The indices that sort each tensor along the batch dimension

Example usage:

>>> 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: Any, descending: bool = False, **kwargs: Any
) -> Any

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	Any	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
Any	The indices that sort each tensor along the sequence dimension.

Example usage:

>>> 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: Any,
    dtype: dtype | None = None,
    device: device | None = None,
) -> Any

Create a new nested data structure with torch.Tensors.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be a torch.Tensor compatible value.	required
dtype	dtype | None	The desired data type of returned tensors. If None, it infers data type from data.	None
device	device | None	The device of the constructed tensors. If None and data is a tensor then the device of data is used. If None and data is not a tensor then the result tensor is constructed on the current device.	None

Returns:

Type	Description
Any	A nested data structure with torch.Tensors. The output data has the same structure as the input.

Example usage:

>>> import numpy as np
>>> from batchtensor.nested import as_tensor
>>> 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)}

batchtensor.nested.asin

asin(data: Any) -> Any

Return new tensors with the arcsine of each element.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

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

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

Return new tensors with the arctangent of each element.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

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

Parameters:

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

Returns:

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

Example usage:

>>> 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

Returns:

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

Example usage:

>>> 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

Returns:

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

Example usage:

>>> 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 chuncks.

Example usage:

>>> 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 chuncks.

Example usage:

>>> 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: Any,
    min: float | None = None,
    max: float | None = None,
) -> Any

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

Parameters:

Name	Type	Description	Default
data	Any	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
Any	The clamp value of the elements. The output has the same structure as the input.

Example usage:

>>> 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.cos

cos(data: Any) -> Any

Return new tensors with the cosine of each element.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

Return new tensors with the hyperbolic cosine of each element.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

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	Any	The input data. Each item must be a tensor.	required

Returns:

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

Example usage:

>>> 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: Any) -> Any

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	Any	The input data. Each item must be a tensor.	required

Returns:

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

Example usage:

>>> 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: Any) -> Any

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	Any	The input data. Each item must be a tensor.	required

Returns:

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

Example usage:

>>> 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: Any) -> Any

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	Any	The input data. Each item must be a tensor.	required

Returns:

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

Example usage:

>>> 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.exp

exp(data: Any) -> Any

Return new tensors with the exponential of the elements.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

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

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

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

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

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

Note

The returned torch.Tensors and numpy.ndarrays share the same memory. Modifications to the torch.Tensors will be reflected in the numpy.ndarrays and vice versa.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be a numpy.ndarray.	required

Returns:

Type	Description
Any	A nested data structure with torch.Tensors instead of numpy.ndarrays. The output data has the same structure as the input.

Example usage:

>>> 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])}

batchtensor.nested.index_select_along_batch

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

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	Any	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
Any	The indexed tensors along the batch dimension.

Example usage:

>>> 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: Any, index: Tensor) -> Any

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	Any	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
Any	The indexed tensors along the sequence dimension.

Example usage:

>>> 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: Any) -> Any

Return new tensors with the natural logarithm of the elements.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

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

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

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

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

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

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	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 usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	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 usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The mean of all elements along the batch dimension.

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

Return the mean of all elements along the sequence dimension.

Note

This function assumes the sequence dimension is the second dimension.

Parameters:

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

Note

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

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	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 usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	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 usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	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 usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	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 usage:

>>> 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.permute_along_batch

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

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	Any	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
Any	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 usage:

>>> 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: Any, permutation: Tensor) -> Any

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	Any	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
Any	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 usage:

>>> 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.prod_along_batch

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

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	Any	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
Any	The product of all elements along the batch dimension.

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The product of all elements along the sequence dimension.

Example usage:

>>> 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: Any, repeats: int) -> Any

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	Any	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
Any	The tensors repeated along the sequence dimension.

Example usage:

>>> 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: Any, index: int) -> Any

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	Any	The input data. Each item must be a tensor.	required
index	int	The index to select with.	required

Returns:

Type	Description
Any	The sliced tensors along the batch dimension.

Example usage:

>>> 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: Any, index: int) -> Any

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	Any	The input data. Each item must be a tensor.	required
index	int	The index to select with.	required

Returns:

Type	Description
Any	The sliced tensors along the sequence dimension.

Example usage:

>>> 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: Any, generator: Generator | None = None
) -> Any

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	Any	The input data. Each item must be a tensor.	required
generator	Generator | None	An optional random number generator.	None

Returns:

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

Example usage:

>>> 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: Any, generator: Generator | None = None
) -> Any

Shuffle all the tensors along the batch 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	Any	The input data. Each item must be a tensor.	required
generator	Generator | None	An optional random number generator.	None

Returns:

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

Example usage:

>>> 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: Any) -> Any

Return new tensors with the sine of each element.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

Return new tensors with the hyperbolic sine of each element.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any,
    start: int = 0,
    stop: int | None = None,
    step: int = 1,
) -> Any

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	Any	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
Any	The sliced tensor along the batch dimension.

Example usage:

>>> 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: Any,
    start: int = 0,
    stop: int | None = None,
    step: int = 1,
) -> Any

Slice all the tensors along the batch 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	Any	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
Any	The sliced tensor along the batch dimension.

Example usage:

>>> 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: Any, descending: bool = False, **kwargs: Any
) -> Any

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	Any	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
Any	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 usage:

>>> 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: Any, descending: bool = False, **kwargs: Any
) -> Any

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	Any	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
Any	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 usage:

>>> 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 chuncks.

Example usage:

>>> 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 chuncks.

Example usage:

>>> 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.sum_along_batch

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

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	Any	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
Any	The sum of all elements along the batch dimension.

Example usage:

>>> 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: Any, keepdim: bool = False) -> Any

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	Any	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
Any	The sum of all elements along the sequence dimension.

Example usage:

>>> 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: Any) -> Any

Return new tensors with the tangent of each element.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any) -> Any

Return new tensors with the hyperbolic tangent of each element.

Parameters:

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

Returns:

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

Example usage:

>>> 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: Any, *args: Any, **kwargs: Any) -> Any

Perform Tensor dtype and/or device conversion.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be a tensor.	required
args	Any	Positional arguments for torch.Tensor.to.	()
kwargs	Any	Keyword arguments for torch.Tensor.to.	{}

Returns:

Type	Description
Any	The data after conversion.

Example usage:

>>> 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]),
... }
>>> 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.])}

batchtensor.nested.to_numpy

to_numpy(data: Any) -> Any

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

Note

The returned torch.Tensors and numpy.ndarrays share the same memory. Modifications to the torch.Tensors will be reflected in the numpy.ndarrays and vice versa.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be a torch.Tensor.	required

Returns:

Type	Description
Any	A nested data structure with numpy.ndarrays instead of torch.Tensors. The output data has the same structure as the input.

Example usage:

>>> import numpy as np
>>> 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])}