nested

batcharray.nested

Contain functions to manipulate nested data.

batcharray.nested.abs

abs(data: Any) -> Any

Return new arrays with the absolute value of each element.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import abs
>>> data = {
...     "a": np.array([[-4, -3], [-2, -1], [0, 1], [2, 3], [4, 5]]),
...     "b": np.array([2, 1, 0, -1, -2]),
... }
>>> out = abs(data)
>>> out
{'a': array([[4, 3], [2, 1], [0, 1], [2, 3], [4, 5]]), 'b': array([2, 1, 0, 1, 2])}

batcharray.nested.amax_along_batch

amax_along_batch(data: Any, keepdims: 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 arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import amax_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = amax_along_batch(data)
>>> out
{'a': array([8, 9]), 'b': np.int64(4)}
>>> out = amax_along_batch(data, keepdims=True)
>>> out
{'a': array([[8, 9]]), 'b': array([4])}

batcharray.nested.amax_along_seq

amax_along_seq(data: Any, keepdims: 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 arrays. All the arrays should have the same sequence size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import amax_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = amax_along_seq(data)
>>> out
{'a': array([4, 9]), 'b': array([4])}
>>> out = amax_along_seq(data, keepdims=True)
>>> out
{'a': array([[4], [9]]), 'b': array([[4]])}

batcharray.nested.amin_along_batch

amin_along_batch(data: Any, keepdims: 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 arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import amin_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = amin_along_batch(data)
>>> out
{'a': array([0, 1]), 'b': np.int64(0)}
>>> out = amin_along_batch(data, keepdims=True)
>>> out
{'a': array([[0, 1]]), 'b': array([0])}

batcharray.nested.amin_along_seq

amin_along_seq(data: Any, keepdims: 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 arrays. All the arrays should have the same sequence size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import amin_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = amin_along_seq(data)
>>> out
{'a': array([0, 5]), 'b': array([0])}
>>> out = amin_along_seq(data, keepdims=True)
>>> out
{'a': array([[0], [5]]), 'b': array([[0]])}

batcharray.nested.arccos

arccos(data: Any) -> Any

Return new arrays with the inverse cosine of each element.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import arccos
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = arccos(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.arccosh

arccosh(data: Any) -> Any

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import arccosh
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = arccosh(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.arcsin

arcsin(data: Any) -> Any

Return new arrays with the arcsine of each element.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import arcsin
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = arcsin(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.arcsinh

arcsinh(data: Any) -> Any

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import arcsinh
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = arcsinh(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.arctan

arctan(data: Any) -> Any

Return new arrays with the arctangent of each element.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import arctan
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = arctan(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.arctanh

arctanh(data: Any) -> Any

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import arctanh
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = arctanh(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.argmax_along_batch

argmax_along_batch(
    data: Any, keepdims: 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 arrays. All the arrays should have the same batch size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be a array.	required
keepdims	bool	Whether the output array 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 numpy as np
>>> from batcharray.nested import argmax_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = argmax_along_batch(data)
>>> out
{'a': array([4, 4]), 'b': np.int64(0)}
>>> out = argmax_along_batch(data, keepdims=True)
>>> out
{'a': array([[4, 4]]), 'b': array([0])}

batcharray.nested.argmax_along_seq

argmax_along_seq(data: Any, keepdims: 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 arrays. All the arrays should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be a array.	required
keepdims	bool	Whether the output array 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 numpy as np
>>> from batcharray.nested import argmax_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = argmax_along_seq(data)
>>> out
{'a': array([4, 4]), 'b': array([0])}
>>> out = argmax_along_seq(data, keepdims=True)
>>> out
{'a': array([[4], [4]]), 'b': array([[0]])}

batcharray.nested.argmin_along_batch

argmin_along_batch(
    data: Any, keepdims: 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 arrays. All the arrays should have the same batch size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be a array.	required
keepdims	bool	Whether the output array 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 numpy as np
>>> from batcharray.nested import argmin_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = argmin_along_batch(data)
>>> out
{'a': array([0, 0]), 'b': np.int64(4)}
>>> out = argmin_along_batch(data, keepdims=True)
>>> out
{'a': array([[0, 0]]), 'b': array([4])}

batcharray.nested.argmin_along_seq

argmin_along_seq(data: Any, keepdims: 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 arrays. All the arrays should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be a array.	required
keepdims	bool	Whether the output array 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 numpy as np
>>> from batcharray.nested import argmin_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = argmin_along_seq(data)
>>> out
{'a': array([0, 0]), 'b': array([4])}
>>> out = argmin_along_seq(data, keepdims=True)
>>> out
{'a': array([[0], [0]]), 'b': array([[4]])}

batcharray.nested.argsort_along_batch

argsort_along_batch(
    data: Any, kind: SortKind | None = None
) -> Any

Return the indices that sort each array along the batch dimension in ascending order by value.

Note

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required
kind	SortKind | None	Sorting algorithm. The default is quicksort. Note that both stable and mergesort use timsort under the covers and, in general, the actual implementation will vary with datatype. The mergesort option is retained for backwards compatibility.	None

Returns:

Type	Description
Any	The indices that sort each array along the batch dimension

Example usage:

>>> import numpy as np
>>> from batcharray.nested import argsort_along_batch
>>> data = {
...     "a": np.array([[2, 6], [0, 3], [4, 9], [8, 1], [5, 7]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = argsort_along_batch(data)
>>> out
{'a': array([[1, 3], [0, 1], [2, 0], [4, 4], [3, 2]]), 'b': array([4, 3, 2, 1, 0])}

batcharray.nested.argsort_along_seq

argsort_along_seq(
    data: Any, kind: SortKind | None = None
) -> Any

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

Note

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required
kind	SortKind | None	Sorting algorithm. The default is quicksort. Note that both stable and mergesort use timsort under the covers and, in general, the actual implementation will vary with datatype. The mergesort option is retained for backwards compatibility.	None

Returns:

Type	Description
Any	The indices that sort each array along the sequence dimension.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import argsort_along_seq
>>> data = {
...     "a": np.array([[7, 3, 0, 8, 5], [1, 9, 6, 4, 2]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = argsort_along_seq(data)
>>> out
{'a': array([[2, 1, 4, 0, 3], [0, 4, 3, 2, 1]]), 'b': array([[4, 3, 2, 1, 0]])}

batcharray.nested.chunk_along_batch

chunk_along_batch(
    data: dict[Hashable, ndarray], chunks: int
) -> list[dict[Hashable, ndarray]]

Split all the arrays into chunks along the batch axis.

Each chunk is a view of the input array.

Note

This function assumes the batch axis is the first axis of the arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

Type	Description
list[dict[Hashable, ndarray]]	The data chuncks.

Example usage:

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

batcharray.nested.chunk_along_seq

chunk_along_seq(
    data: dict[Hashable, ndarray], chunks: int
) -> list[dict[Hashable, ndarray]]

Split all the arrays into chunks along the sequence axis.

Each chunk is a view of the input array.

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

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

Returns:

Type	Description
list[dict[Hashable, ndarray]]	The data chuncks.

Example usage:

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

batcharray.nested.clip

clip(
    data: Any,
    a_min: float | None = None,
    a_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 an array .	required
a_min	float | None	The lower-bound of the range to be clamped to.	None
a_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 numpy as np
>>> from batcharray.nested import clip
>>> data = {
...     "a": np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": np.array([5, 4, 3, 2, 1]),
... }
>>> out = clip(data, a_min=1, a_max=5)
>>> out
{'a': array([[1, 2], [3, 4], [5, 5], [5, 5], [5, 5]]), 'b': array([5, 4, 3, 2, 1])}

batcharray.nested.concatenate_along_batch

concatenate_along_batch(
    data: Sequence[dict[Hashable, ndarray]]
) -> dict[Hashable, ndarray]

Concatenate the given arrays in the batch axis.

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

Note

This function assumes the batch axis is the first axis of the arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

Type	Description
dict[Hashable, ndarray]	The concatenated arrays along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import concatenate_along_batch
>>> data = [
...     {
...         "a": np.array([[0, 1, 2], [4, 5, 6]]),
...         "b": np.array([[10, 11, 12], [13, 14, 15]]),
...     },
...     {"a": np.array([[7, 8, 9]]), "b": np.array([[17, 18, 19]])},
... ]
>>> out = concatenate_along_batch(data)
>>> out
{'a': array([[0, 1, 2], [4, 5, 6], [7, 8, 9]]),
 'b': array([[10, 11, 12], [13, 14, 15], [17, 18, 19]])}

batcharray.nested.concatenate_along_seq

concatenate_along_seq(
    data: Sequence[dict[Hashable, ndarray]]
) -> dict[Hashable, ndarray]

Concatenate the given arrays in the sequence axis.

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

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

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

Returns:

Type	Description
dict[Hashable, ndarray]	The concatenated arrays along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import concatenate_along_seq
>>> data = [
...     {
...         "a": np.array([[0, 1, 2], [4, 5, 6]]),
...         "b": np.array([[10, 11, 12], [13, 14, 15]]),
...     },
...     {"a": np.array([[7], [8]]), "b": np.array([[17], [18]])},
... ]
>>> out = concatenate_along_seq(data)
>>> out
{'a': array([[0, 1, 2, 7], [4, 5, 6, 8]]),
 'b': array([[10, 11, 12, 17], [13, 14, 15, 18]])}

batcharray.nested.cos

cos(data: Any) -> Any

Return new arrays with the cosine of each element.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import cos
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = cos(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.cosh

cosh(data: Any) -> Any

Return new arrays with the hyperbolic cosine of each element.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import cosh
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = cosh(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.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 axis is the first axis of the arrays. All the arrays should have the same batch size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required

Returns:

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

Example usage:

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

batcharray.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 axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required

Returns:

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

Example usage:

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

batcharray.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 axis is the first axis of the arrays. All the arrays should have the same batch size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required

Returns:

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

Example usage:

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

batcharray.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 axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required

Returns:

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

Example usage:

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

batcharray.nested.exp

exp(data: Any) -> Any

Return new arrays with the exponential of the elements.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import exp
>>> data = {
...     "a": np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": np.array([5, 4, 3, 2, 1]),
... }
>>> out = exp(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.exp2

exp2(data: Any) -> Any

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import exp2
>>> data = {
...     "a": np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": np.array([5, 4, 3, 2, 1]),
... }
>>> out = exp2(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.expm1

expm1(data: Any) -> Any

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import expm1
>>> data = {
...     "a": np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": np.array([5, 4, 3, 2, 1]),
... }
>>> out = expm1(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.index_select_along_batch

index_select_along_batch(
    data: Any, indices: ndarray
) -> Any

Return the arrays which index the arrays along the batch axis using the entries in indices.

Note

This function assumes the batch axis is the first axis of the arrays. All the arrays should have the same batch size.

Note

Equivalent to take_along_batch.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
indices	ndarray	The 1-D array containing the indices to take.	required

Returns:

Type	Description
Any	The indexed arrays along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import index_select_along_batch
>>> arrays = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = index_select_along_batch(arrays, np.array([2, 4]))
>>> out
{'a': array([[4, 5], [8, 9]]), 'b': array([2, 0])}
>>> out = index_select_along_batch(arrays, np.array([4, 3, 2, 1, 0]))
>>> out
{'a': array([[8, 9], [6, 7], [4, 5], [2, 3], [0, 1]]), 'b': array([0, 1, 2, 3, 4])}

batcharray.nested.index_select_along_seq

index_select_along_seq(data: Any, indices: ndarray) -> Any

Return the arrays which index the arrays along the sequence axis using the entries in indices.

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Note

Equivalent to take_along_seq.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
indices	ndarray	The 1-D array containing the indices to take.	required

Returns:

Type	Description
Any	The indexed arrays along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import index_select_along_seq
>>> arrays = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = index_select_along_seq(arrays, np.array([2, 4]))
>>> out
{'a': array([[2, 4], [7, 9]]), 'b': array([[2, 0]])}
>>> out = index_select_along_seq(arrays, np.array([4, 3, 2, 1, 0]))
>>> out
{'a': array([[4, 3, 2, 1, 0], [9, 8, 7, 6, 5]]), 'b': array([[0, 1, 2, 3, 4]])}

batcharray.nested.log

log(data: Any) -> Any

Return new arrays with the natural logarithm of the elements.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import log
>>> data = {
...     "a": np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": np.array([5, 4, 3, 2, 1]),
... }
>>> out = log(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.log10

log10(data: Any) -> Any

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	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 numpy as np
>>> from batcharray.nested import log10
>>> data = {
...     "a": np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": np.array([5, 4, 3, 2, 1]),
... }
>>> out = log10(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.log1p

log1p(data: Any) -> Any

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import log1p
>>> data = {
...     "a": np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": np.array([5, 4, 3, 2, 1]),
... }
>>> out = log1p(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.log2

log2(data: Any) -> Any

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	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 numpy as np
>>> from batcharray.nested import log2
>>> data = {
...     "a": np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
...     "b": np.array([5, 4, 3, 2, 1]),
... }
>>> out = log2(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.masked_select_along_batch

masked_select_along_batch(data: Any, mask: ndarray) -> Any

Return the arrays which index the arrays along the batch axis according to the boolean mask mask.

Note

This function assumes the batch axis is the first axis of the arrays. All the arrays should have the same batch size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
mask	ndarray	The 1-D array containing the binary mask to index with.	required

Returns:

Type	Description
Any	The indexed arrays along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import masked_select_along_batch
>>> arrays = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = masked_select_along_batch(arrays, np.array([False, False, True, False, True]))
>>> out
{'a': array([[4, 5], [8, 9]]), 'b': array([2, 0])}

batcharray.nested.masked_select_along_seq

masked_select_along_seq(data: Any, mask: ndarray) -> Any

Return the arrays which index the arrays along the sequence axis according to the boolean mask mask.

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
mask	ndarray	The 1-D array containing the binary mask to index with.	required

Returns:

Type	Description
Any	The indexed arrays along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import masked_select_along_seq
>>> arrays = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = masked_select_along_seq(arrays, np.array([False, False, True, False, True]))
>>> out
{'a': array([[2, 4], [7, 9]]), 'b': array([[2, 0]])}

batcharray.nested.max_along_batch

max_along_batch(data: Any, keepdims: 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 arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import max_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = max_along_batch(data)
>>> out
{'a': array([8, 9]), 'b': np.int64(4)}
>>> out = max_along_batch(data, keepdims=True)
>>> out
{'a': array([[8, 9]]), 'b': array([4])}

batcharray.nested.max_along_seq

max_along_seq(data: Any, keepdims: 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 arrays. All the arrays should have the same sequence size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import max_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = max_along_seq(data)
>>> out
{'a': array([4, 9]), 'b': array([4])}
>>> out = max_along_seq(data, keepdims=True)
>>> out
{'a': array([[4], [9]]), 'b': array([[4]])}

batcharray.nested.mean_along_batch

mean_along_batch(data: Any, keepdims: 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 arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import mean_along_batch
>>> data = {
...     "a": np.array([[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0], [8.0, 9.0]]),
...     "b": np.array([4, 3, 2, 1, 0], dtype=np.float32),
... }
>>> out = mean_along_batch(data)
>>> out
{'a': array([4., 5.]), 'b': np.float32(2.0)}
>>> out = mean_along_batch(data, keepdims=True)
>>> out
{'a': array([[4., 5.]]), 'b': array([2.], dtype=float32)}

batcharray.nested.mean_along_seq

mean_along_seq(data: Any, keepdims: 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 array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Note

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import mean_along_seq
>>> data = {
...     "a": np.array([[0.0, 1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0, 9.0]]),
...     "b": np.array([[4, 3, 2, 1, 0]], dtype=np.float32),
... }
>>> out = mean_along_seq(data)
>>> out
{'a': array([2., 7.]), 'b': array([2.], dtype=float32)}
>>> out = mean_along_seq(data, keepdims=True)
>>> out
{'a': array([[2.], [7.]]), 'b': array([[2.]], dtype=float32)}

batcharray.nested.median_along_batch

median_along_batch(
    data: Any, keepdims: 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 arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import median_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = median_along_batch(data)
>>> out
{'a': array([4., 5.]), 'b': np.float64(2.0)}
>>> out = median_along_batch(data, keepdims=True)
>>> out
{'a': array([[4., 5.]]), 'b': array([2.])}

batcharray.nested.median_along_seq

median_along_seq(data: Any, keepdims: 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 arrays. All the arrays should have the same sequence size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import median_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = median_along_seq(data)
>>> out
{'a': array([2., 7.]), 'b': array([2.])}
>>> out = median_along_seq(data, keepdims=True)
>>> out
{'a': array([[2.], [7.]]), 'b': array([[2.]])}

batcharray.nested.min_along_batch

min_along_batch(data: Any, keepdims: 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 arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import min_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = min_along_batch(data)
>>> out
{'a': array([0, 1]), 'b': np.int64(0)}
>>> out = min_along_batch(data, keepdims=True)
>>> out
{'a': array([[0, 1]]), 'b': array([0])}

batcharray.nested.min_along_seq

min_along_seq(data: Any, keepdims: 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 arrays. All the arrays should have the same sequence size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import min_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = min_along_seq(data)
>>> out
{'a': array([0, 5]), 'b': array([0])}
>>> out = min_along_seq(data, keepdims=True)
>>> out
{'a': array([[0], [5]]), 'b': array([[0]])}

batcharray.nested.permute_along_batch

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

Permute all the arrays along the batch axis.

Note

This function assumes the batch axis is the first axis of the arrays. All the arrays should have the same batch size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
permutation	ndarray	The 1-D array containing the indices of the permutation. The shape should match the batch axis of the array.	required

Returns:

Type	Description
Any	The data with permuted arrays along the batch axis. 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 axis of the array.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import permute_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = permute_along_batch(data, np.array([2, 1, 3, 0, 4]))
>>> out
{'a': array([[4, 5], [2, 3], [6, 7], [0, 1], [8, 9]]), 'b': array([2, 3, 1, 4, 0])}

batcharray.nested.permute_along_seq

permute_along_seq(data: Any, permutation: ndarray) -> Any

Permute all the arrays along the sequence axis.

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
permutation	ndarray	The 1-D array containing the indices of the permutation. The shape should match the sequence axis of the array.	required

Returns:

Type	Description
Any	The data with permuted arrays along the sequence axis. 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 axis of the array.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import permute_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = permute_along_seq(data, np.array([2, 1, 3, 0, 4]))
>>> out
{'a': array([[2, 1, 3, 0, 4], [7, 6, 8, 5, 9]]), 'b': array([[2, 3, 1, 4, 0]])}

batcharray.nested.prod_along_batch

prod_along_batch(data: Any, keepdims: 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 arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import prod_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([5, 4, 3, 2, 1]),
... }
>>> out = prod_along_batch(data)
>>> out
{'a': array([  0, 945]), 'b': np.int64(120)}
>>> out = prod_along_batch(data, keepdims=True)
>>> out
{'a': array([[  0, 945]]), 'b': array([120])}

batcharray.nested.prod_along_seq

prod_along_seq(data: Any, keepdims: 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 arrays. All the arrays should have the same sequence size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import prod_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[5, 4, 3, 2, 1]]),
... }
>>> out = prod_along_seq(data)
>>> out
{'a': array([    0, 15120]), 'b': array([120])}
>>> out = prod_along_seq(data, keepdims=True)
>>> out
{'a': array([[    0], [15120]]), 'b': array([[120]])}

batcharray.nested.select_along_batch

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

Slice the arrays along the batch axis at the given index.

This function returns a view of the original array with the batch axis removed.

Note

This function assumes the batch axis is the first axis of the arrays. All the arrays should have the same batch size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
index	int	The index to select with.	required

Returns:

Type	Description
Any	The sliced arrays along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import select_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = select_along_batch(data, index=2)
>>> out
{'a': array([4, 5]), 'b': np.int64(2)}

batcharray.nested.select_along_seq

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

Slice the arrays along the sequence axis at the given index.

This function returns a view of the original array with the sequence axis removed.

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
index	int	The index to select with.	required

Returns:

Type	Description
Any	The sliced arrays along the sequence axis.

Example usage:

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

batcharray.nested.shuffle_along_batch

shuffle_along_batch(
    data: Any, rng: Generator | None = None
) -> Any

Shuffle all the arrays along the batch axis.

Note

This function assumes the batch axis is the first axis of the arrays. All the arrays should have the same batch size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
rng	Generator | None	An optional random number generator.	None

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import shuffle_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = shuffle_along_batch(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.shuffle_along_seq

shuffle_along_seq(
    data: Any, rng: Generator | None = None
) -> Any

Shuffle all the arrays along the batch axis.

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
rng	Generator | None	An optional random number generator.	None

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import shuffle_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = shuffle_along_seq(data)
>>> out
{'a': array([[...]]), 'b': array([[...]])}

batcharray.nested.sin

sin(data: Any) -> Any

Return new arrays with the sine of each element.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import sin
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = sin(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.sinh

sinh(data: Any) -> Any

Return new arrays with the hyperbolic sine of each element.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import sinh
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = sinh(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.slice_along_batch

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

Slice all the arrays along the batch axis.

Note

This function assumes the batch axis is the first axis of the arrays. All the arrays should have the same batch size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	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 array along the batch axis.

Example usage:

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

batcharray.nested.slice_along_seq

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

Slice all the arrays along the batch axis.

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	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 array along the batch axis.

Example usage:

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

batcharray.nested.sort_along_batch

sort_along_batch(
    data: Any, kind: SortKind | None = None
) -> Any

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

Note

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required
kind	SortKind | None	Sorting algorithm. The default is quicksort. Note that both stable and mergesort use timsort under the covers and, in general, the actual implementation will vary with datatype. The mergesort option is retained for backwards compatibility.	None

Returns:

Type	Description
Any	A similar object where each array is replaced by a sorted array along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import sort_along_batch
>>> data = {
...     "a": np.array([[2, 6], [0, 3], [4, 9], [8, 1], [5, 7]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = sort_along_batch(data)
>>> out
{'a': array([[0, 1], [2, 3], [4, 6], [5, 7], [8, 9]]), 'b': array([0, 1, 2, 3, 4])}

batcharray.nested.sort_along_seq

sort_along_seq(
    data: Any, kind: SortKind | None = None
) -> Any

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

Note

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

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required
kind	SortKind | None	Sorting algorithm. The default is quicksort. Note that both stable and mergesort use timsort under the covers and, in general, the actual implementation will vary with datatype. The mergesort option is retained for backwards compatibility.	None

Returns:

Type	Description
Any	A similar object where each array is replaced by a sorted array along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import sort_along_seq
>>> data = {
...     "a": np.array([[7, 3, 0, 8, 5], [1, 9, 6, 4, 2]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = sort_along_seq(data)
>>> out
{'a': array([[0, 3, 5, 7, 8], [1, 2, 4, 6, 9]]), 'b': array([[0, 1, 2, 3, 4]])}

batcharray.nested.split_along_batch

split_along_batch(
    data: dict[Hashable, ndarray],
    split_size_or_sections: int | Sequence[int],
) -> list[dict[Hashable, ndarray]]

Split all the arrays into chunks along the batch axis.

Each chunk is a view of the original array.

Note

This function assumes the batch axis is the first axis of the arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

Type	Description
list[dict[Hashable, ndarray]]	The data chuncks.

Example usage:

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

batcharray.nested.split_along_seq

split_along_seq(
    data: dict[Hashable, ndarray],
    split_size_or_sections: int | Sequence[int],
) -> list[dict[Hashable, ndarray]]

Split all the arrays into chunks along the sequence axis.

Each chunk is a view of the original array.

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

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

Returns:

Type	Description
list[dict[Hashable, ndarray]]	The data chuncks.

Example usage:

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

batcharray.nested.sum_along_batch

sum_along_batch(data: Any, keepdims: 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 arrays. All the arrays should have the same batch size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import sum_along_batch
>>> data = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = sum_along_batch(data)
>>> out
{'a': array([20, 25]), 'b': np.int64(10)}
>>> out = sum_along_batch(data, keepdims=True)
>>> out
{'a': array([[20, 25]]), 'b': array([10])}

batcharray.nested.sum_along_seq

sum_along_seq(data: Any, keepdims: 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 arrays. All the arrays should have the same sequence size.

Parameters:

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

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import sum_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = sum_along_seq(data)
>>> out
{'a': array([10, 35]), 'b': array([10])}
>>> out = sum_along_seq(data, keepdims=True)
>>> out
{'a': array([[10], [35]]), 'b': array([[10]])}

batcharray.nested.take_along_batch

take_along_batch(data: Any, indices: ndarray) -> Any

Return the arrays which index the arrays along the batch axis using the entries in indices.

Note

This function assumes the batch axis is the first axis of the arrays. All the arrays should have the same batch size.

Note

Equivalent to index_select_along_batch.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
indices	ndarray	The 1-D array containing the indices to take.	required

Returns:

Type	Description
Any	The indexed arrays along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import take_along_batch
>>> arrays = {
...     "a": np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]),
...     "b": np.array([4, 3, 2, 1, 0]),
... }
>>> out = take_along_batch(arrays, np.array([2, 4]))
>>> out
{'a': array([[4, 5], [8, 9]]), 'b': array([2, 0])}
>>> out = take_along_batch(arrays, np.array([4, 3, 2, 1, 0]))
>>> out
{'a': array([[8, 9], [6, 7], [4, 5], [2, 3], [0, 1]]), 'b': array([0, 1, 2, 3, 4])}

batcharray.nested.take_along_seq

take_along_seq(data: Any, indices: ndarray) -> Any

Return the arrays which index the arrays along the sequence axis using the entries in indices.

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Note

Equivalent to index_select_along_seq.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
indices	ndarray	The 1-D array containing the indices to take.	required

Returns:

Type	Description
Any	The indexed arrays along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import take_along_seq
>>> arrays = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = take_along_seq(arrays, np.array([2, 4]))
>>> out
{'a': array([[2, 4], [7, 9]]), 'b': array([[2, 0]])}
>>> out = take_along_seq(arrays, np.array([4, 3, 2, 1, 0]))
>>> out
{'a': array([[4, 3, 2, 1, 0], [9, 8, 7, 6, 5]]), 'b': array([[0, 1, 2, 3, 4]])}

batcharray.nested.tan

tan(data: Any) -> Any

Return new arrays with the tangent of each element.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import tan
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = tan(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.tanh

tanh(data: Any) -> Any

Return new arrays with the hyperbolic tangent of each element.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array .	required

Returns:

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

Example usage:

>>> import numpy as np
>>> from batcharray.nested import tanh
>>> data = {"a": np.ones((2, 3)), "b": np.arange(5)}
>>> out = tanh(data)
>>> out
{'a': array([[...]]), 'b': array([...])}

batcharray.nested.tile_along_seq

tile_along_seq(data: Any, reps: int) -> Any

Repeat all the arrays along the sequence axis.

Note

This function assumes the sequence axis is the second axis of the arrays. All the arrays should have the same sequence size.

Parameters:

Name	Type	Description	Default
data	Any	The input data. Each item must be an array.	required
reps	int	The number of repetitions data along the sequence axis.	required

Returns:

Type	Description
Any	The arrays repeated along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import tile_along_seq
>>> data = {
...     "a": np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]),
...     "b": np.array([[4, 3, 2, 1, 0]]),
... }
>>> out = tile_along_seq(data, 2)
>>> out
{'a': array([[0, 1, 2, 3, 4, 0, 1, 2, 3, 4], [5, 6, 7, 8, 9, 5, 6, 7, 8, 9]]),
 'b': array([[4, 3, 2, 1, 0, 4, 3, 2, 1, 0]])}

batcharray.nested.to_list

to_list(data: Any) -> Any

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

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 `lists instead of numpy.ndarrays. The output data has the same structure as the input.

Example usage:

>>> import numpy as np
>>> from batcharray.nested import to_list
>>> data = {"a": np.ones((2, 5)), "b": np.array([0, 1, 2, 3, 4])}
>>> out = to_list(data)
>>> out
{'a': [[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]], 'b': [0, 1, 2, 3, 4]}