array

batcharray.array

Contain functions to manipulate arrays.

batcharray.array.amax_along_batch

amax_along_batch(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the maximum of all elements along the batch axis.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The maximum of all elements along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import amax_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = amax_along_batch(array)
>>> out
array([8, 9])
>>> out = amax_along_batch(array, keepdims=True)
>>> out
array([[8, 9]])

batcharray.array.amax_along_seq

amax_along_seq(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the maximum of all elements along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The maximum of all elements along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import amax_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = amax_along_seq(array)
>>> out
array([4, 9])
>>> out = amax_along_seq(array, keepdims=True)
>>> out
array([[4], [9]])

batcharray.array.amin_along_batch

amin_along_batch(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the minimum of all elements along the batch axis.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The minimum of all elements along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import amin_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = amin_along_batch(array)
>>> out
array([0, 1])
>>> out = amin_along_batch(array, keepdims=True)
>>> out
array([[0, 1]])

batcharray.array.amin_along_seq

amin_along_seq(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the minimum of all elements along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The minimum of all elements along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import amin_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = amin_along_seq(array)
>>> out
array([0, 5])
>>> out = amin_along_seq(array, keepdims=True)
>>> out
array([[0], [5]])

batcharray.array.argmax_along_batch

argmax_along_batch(
    array: ndarray, keepdims: bool = False
) -> ndarray

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

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The indices of the maximum value of all elements along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import argmax_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = argmax_along_batch(array)
>>> out
array([4, 4])
>>> out = argmax_along_batch(array, keepdims=True)
>>> out
array([[4, 4]])

batcharray.array.argmax_along_seq

argmax_along_seq(
    array: ndarray, keepdims: bool = False
) -> ndarray

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

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The indices of the maximum value of all elements along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import argmax_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = argmax_along_seq(array)
>>> out
array([4, 4])
>>> out = argmax_along_seq(array, keepdims=True)
>>> out
array([[4], [4]])

batcharray.array.argmin_along_batch

argmin_along_batch(
    array: ndarray, keepdims: bool = False
) -> ndarray

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

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The indices of the minimum value of all elements along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import argmin_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = argmin_along_batch(array)
>>> out
array([0, 0])
>>> out = argmin_along_batch(array, keepdims=True)
>>> out
array([[0, 0]])

batcharray.array.argmin_along_seq

argmin_along_seq(
    array: ndarray, keepdims: bool = False
) -> ndarray

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

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The indices of the minimum value of all elements along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import argmin_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = argmin_along_seq(array)
>>> out
array([0, 0])
>>> out = argmin_along_seq(array, keepdims=True)
>>> out
array([[0], [0]])

batcharray.array.argsort_along_batch

argsort_along_batch(
    array: ndarray, kind: SortKind | None = None
) -> ndarray

Return the indices that sort an array along the batch axis in ascending order by value.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input 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
ndarray	The indices that sort the array along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import argsort_along_batch
>>> array = np.array([[2, 6], [0, 3], [4, 9], [8, 1], [5, 7]])
>>> out = argsort_along_batch(array)
>>> out
array([[1, 3], [0, 1], [2, 0], [4, 4], [3, 2]])

batcharray.array.argsort_along_seq

argsort_along_seq(
    array: ndarray, kind: SortKind | None = None
) -> ndarray

Return the indices that sort an array along the sequence axis in ascending order by value.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input 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
ndarray	The indices that sort the array along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import argsort_along_seq
>>> array = np.array([[7, 3, 0, 8, 5], [1, 9, 6, 4, 2]])
>>> out = argsort_along_seq(array)
>>> out
array([[2, 1, 4, 0, 3],
       [0, 4, 3, 2, 1]])

batcharray.array.chunk_along_batch

chunk_along_batch(
    array: ndarray, chunks: int
) -> list[ndarray]

Split the array 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.

Parameters:

Name	Type	Description	Default
array	ndarray	The array to split.	required
chunks	int	Number of chunks to return.	required

Returns:

Type	Description
list[ndarray]	The array chunks.

Example usage:

>>> import numpy as np
>>> from batcharray.array import chunk_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> outputs = chunk_along_batch(array, chunks=3)
>>> outputs
[array([[0, 1], [2, 3]]), array([[4, 5], [6, 7]]), array([[8, 9]])]

batcharray.array.chunk_along_seq

chunk_along_seq(
    array: ndarray, chunks: int
) -> list[ndarray]

Split the array 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.

Parameters:

Name	Type	Description	Default
array	ndarray	The array to split.	required
chunks	int	Number of chunks to return.	required

Returns:

Type	Description
list[ndarray]	The array chunks.

Example usage:

>>> import numpy as np
>>> from batcharray.array import chunk_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> outputs = chunk_along_seq(array, chunks=3)
>>> outputs
[array([[0, 1], [5, 6]]), array([[2, 3], [7, 8]]), array([[4], [9]])]

batcharray.array.concatenate_along_batch

concatenate_along_batch(
    arrays: list[ndarray] | tuple[ndarray, ...],
) -> 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.

Parameters:

Name	Type	Description	Default
arrays	list[ndarray] | tuple[ndarray, ...]	The arrays to concatenate.	required

Returns:

Type	Description
ndarray	The concatenated arrays along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import concatenate_along_batch
>>> arrays = [
...     np.array([[0, 1, 2], [4, 5, 6]]),
...     np.array([[10, 11, 12], [13, 14, 15]]),
... ]
>>> out = concatenate_along_batch(arrays)
>>> out
array([[ 0,  1,  2],
       [ 4,  5,  6],
       [10, 11, 12],
       [13, 14, 15]])

batcharray.array.concatenate_along_seq

concatenate_along_seq(
    arrays: list[ndarray] | tuple[ndarray, ...],
) -> 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.

Parameters:

Name	Type	Description	Default
arrays	list[ndarray] | tuple[ndarray, ...]	The arrays to concatenate.	required

Returns:

Type	Description
ndarray	The concatenated arrays along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import concatenate_along_seq
>>> arrays = [
...     np.array([[0, 1, 2], [4, 5, 6]]),
...     np.array([[10, 11], [12, 13]]),
... ]
>>> out = concatenate_along_seq(arrays)
>>> out
array([[ 0,  1,  2, 10, 11],
       [ 4,  5,  6, 12, 13]])

batcharray.array.cumprod_along_batch

cumprod_along_batch(array: ndarray) -> ndarray

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

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required

Returns:

Type	Description
ndarray	The cumulative product of elements of input in the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import cumprod_along_batch
>>> array = np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]])
>>> out = cumprod_along_batch(array)
>>> out
array([[   1,    2], [   3,    8], [  15,   48], [ 105,  384], [ 945, 3840]])

batcharray.array.cumprod_along_seq

cumprod_along_seq(array: ndarray) -> ndarray

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

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required

Returns:

Type	Description
ndarray	The cumulative product of elements of input in the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import cumprod_along_seq
>>> array = np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
>>> out = cumprod_along_seq(array)
>>> out
array([[    1,     2,     6,    24,   120],
        [    6,    42,   336,  3024, 30240]])

batcharray.array.cumsum_along_batch

cumsum_along_batch(array: ndarray) -> ndarray

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

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required

Returns:

Type	Description
ndarray	The cumulative sum of elements of input in the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import cumsum_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = cumsum_along_batch(array)
>>> out
array([[ 0,  1], [ 2,  4], [ 6,  9], [12, 16], [20, 25]])

batcharray.array.cumsum_along_seq

cumsum_along_seq(array: ndarray) -> ndarray

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

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required

Returns:

Type	Description
ndarray	The cumulative sum of elements of input in the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import cumsum_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = cumsum_along_seq(array)
>>> out
array([[ 0,  1,  3,  6, 10],
       [ 5, 11, 18, 26, 35]])

batcharray.array.index_select_along_batch

index_select_along_batch(
    array: ndarray, indices: ndarray
) -> ndarray

Return a new array which indexes the input array along the batch axis using the entries in indices.

Note

This function assumes the batch axis is the first axis.

Note

Equivalent to take_along_batch.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
indices	ndarray	The 1-D array containing the indices to take.	required

Returns:

Type	Description
ndarray	The indexed array along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import index_select_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = index_select_along_batch(array, np.array([2, 4]))
>>> out
array([[4, 5],
       [8, 9]])
>>> out = index_select_along_batch(array, np.array([4, 3, 2, 1, 0]))
>>> out
array([[8, 9],
       [6, 7],
       [4, 5],
       [2, 3],
       [0, 1]])

batcharray.array.index_select_along_seq

index_select_along_seq(
    array: ndarray, indices: ndarray
) -> ndarray

Return a new array which indexes the input array along the sequence axis using the entries in indices.

Note

This function assumes the sequence axis is the second axis.

Note

Equivalent to take_along_seq.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
indices	ndarray	The 1-D array containing the indices to take.	required

Returns:

Type	Description
ndarray	The indexed array along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import index_select_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = index_select_along_seq(array, np.array([2, 4]))
>>> out
array([[2, 4],
       [7, 9]])
>>> out = index_select_along_seq(array, np.array([4, 3, 2, 1, 0]))
>>> out
array([[4, 3, 2, 1, 0],
       [9, 8, 7, 6, 5]])

batcharray.array.masked_select_along_batch

masked_select_along_batch(
    array: ndarray, mask: ndarray
) -> ndarray

Return a new array which indexes the input array along the batch axis according to the boolean mask mask.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
mask	ndarray	The 1-D array containing the binary mask to index with.	required

Returns:

Type	Description
ndarray	The indexed array along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import masked_select_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = masked_select_along_batch(array, np.array([False, False, True, False, True]))
>>> out
array([[4, 5], [8, 9]])

batcharray.array.masked_select_along_seq

masked_select_along_seq(
    array: ndarray, mask: ndarray
) -> ndarray

Return a new array which indexes the input array along the sequence axis according to the boolean mask mask.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
mask	ndarray	The 1-D array containing the binary mask to index with.	required

Returns:

Type	Description
ndarray	The indexed array along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import masked_select_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = masked_select_along_seq(array, np.array([False, False, True, False, True]))
>>> out
array([[2, 4], [7, 9]])

batcharray.array.max_along_batch

max_along_batch(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the maximum of all elements along the batch axis.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The maximum of the input array along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import max_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = max_along_batch(array)
>>> out
array([8, 9])
>>> out = max_along_batch(array, keepdims=True)
>>> out
array([[8, 9]])

batcharray.array.max_along_seq

max_along_seq(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the maximum of all elements along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The minimum of the input array along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import max_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = max_along_seq(array)
>>> out
array([4, 9])
>>> out = max_along_seq(array, keepdims=True)
>>> out
array([[4], [9]])

batcharray.array.mean_along_batch

mean_along_batch(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the mean of all elements along the batch axis.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The mean of all elements along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import mean_along_batch
>>> array = np.array([[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0], [8.0, 9.0]])
>>> out = mean_along_batch(array)
>>> out
array([4., 5.])
>>> out = mean_along_batch(array, keepdims=True)
>>> out
array([[4., 5.]])

batcharray.array.mean_along_seq

mean_along_seq(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the mean of all elements along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The mean of all elements along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import mean_along_seq
>>> array = np.array([[0.0, 1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0, 9.0]])
>>> out = mean_along_seq(array)
>>> out
array([2., 7.])
>>> out = mean_along_seq(array, keepdims=True)
>>> out
array([[2.], [7.]])

batcharray.array.median_along_batch

median_along_batch(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the median of all elements along the batch axis.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	A new array holding the result. If the input contains integers or floats smaller than np.float64, then the output data-type is np.float64. Otherwise, the data-type of the output is the same as that of the input.

Example usage:

>>> import numpy as np
>>> from batcharray.array import median_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = median_along_batch(array)
>>> out
array([4., 5.])
>>> out = median_along_batch(array, keepdims=True)
>>> out
array([[4., 5.]])

batcharray.array.median_along_seq

median_along_seq(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the median of all elements along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	A new array holding the result. If the input contains integers or floats smaller than np.float64, then the output data-type is np.float64. Otherwise, the data-type of the output is the same as that of the input.

Example usage:

>>> import numpy as np
>>> from batcharray.array import median_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = median_along_seq(array)
>>> out
array([2., 7.])
>>> out = median_along_seq(array, keepdims=True)
>>> out
array([[2.], [7.]])

batcharray.array.min_along_batch

min_along_batch(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the minimum of all elements along the batch axis.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The minimum of the input array along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import min_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = min_along_batch(array)
>>> out
array([0, 1])
>>> out = min_along_batch(array, keepdims=True)
>>> out
array([[0, 1]])

batcharray.array.min_along_seq

min_along_seq(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the minimum of all elements along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The minimum of the input array along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import min_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = min_along_seq(array)
>>> out
array([0, 5])
>>> out = min_along_seq(array, keepdims=True)
>>> out
array([[0], [5]])

batcharray.array.permute_along_batch

permute_along_batch(
    array: ndarray, permutation: ndarray
) -> ndarray

Permute the array along the batch axis.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The array to split.	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
ndarray	The array with permuted data along the batch axis.

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.array import permute_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = permute_along_batch(array, np.array([2, 1, 3, 0, 4]))
>>> out
array([[4, 5],
       [2, 3],
       [6, 7],
       [0, 1],
       [8, 9]])

batcharray.array.permute_along_seq

permute_along_seq(
    array: ndarray, permutation: ndarray
) -> ndarray

Permute the array along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The array to split.	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
ndarray	The array with permuted data along the sequence axis.

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.array import permute_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = permute_along_seq(array, np.array([2, 1, 3, 0, 4]))
>>> out
array([[2, 1, 3, 0, 4],
       [7, 6, 8, 5, 9]])

batcharray.array.prod_along_batch

prod_along_batch(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the product of all elements along the batch axis.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The product of all elements along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import prod_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = prod_along_batch(array)
>>> out
array([  0, 945])
>>> out = prod_along_batch(array, keepdims=True)
>>> out
array([[  0, 945]])

batcharray.array.prod_along_seq

prod_along_seq(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the product of all elements along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The product of all elements along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import prod_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = prod_along_seq(array)
>>> out
array([    0, 15120])
>>> out = prod_along_seq(array, keepdims=True)
>>> out
array([[    0], [15120]])

batcharray.array.select_along_batch

select_along_batch(array: ndarray, index: int) -> ndarray

Slice the input array 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.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
index	int	The index to select with.	required

Returns:

Type	Description
ndarray	The sliced array along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import select_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = select_along_batch(array, index=2)
>>> out
array([4, 5])

batcharray.array.select_along_seq

select_along_seq(array: ndarray, index: int) -> ndarray

Slice the input array 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.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
index	int	The index to select with.	required

Returns:

Type	Description
ndarray	The sliced array along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import select_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = select_along_seq(array, index=2)
>>> out
array([2, 7])

batcharray.array.shuffle_along_batch

shuffle_along_batch(
    array: ndarray, rng: Generator | None = None
) -> ndarray

Shuffle the array along the batch dimension.

Note

This function assumes the batch axis is the first dimension.

Parameters:

Name	Type	Description	Default
array	ndarray	The array to split.	required
rng	Generator | None	An optional random number generator.	None

Returns:

Type	Description
ndarray	The shuffled array.

Example usage:

>>> import numpy as np
>>> from batcharray.array import shuffle_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = shuffle_along_batch(array)
>>> out
array([[...]])

batcharray.array.shuffle_along_seq

shuffle_along_seq(
    array: ndarray, rng: Generator | None = None
) -> ndarray

Shuffle the array along the batch dimension.

Note

This function assumes the sequence axis is the second dimension.

Parameters:

Name	Type	Description	Default
array	ndarray	The array to split.	required
rng	Generator | None	An optional random number generator.	None

Returns:

Type	Description
ndarray	The shuffled array.

Example usage:

>>> import numpy as np
>>> from batcharray.array import shuffle_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = shuffle_along_seq(array)
>>> out
array([[...]])

batcharray.array.slice_along_batch

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

Slice the array along the batch axis.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input 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
ndarray	The sliced array along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import slice_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = slice_along_batch(array, start=2)
>>> out
array([[4, 5], [6, 7], [8, 9]])
>>> out = slice_along_batch(array, stop=3)
>>> out
array([[0, 1], [2, 3], [4, 5]])
>>> out = slice_along_batch(array, step=2)
>>> out
array([[0, 1], [4, 5], [8, 9]])

batcharray.array.slice_along_seq

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

Slice the array along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input 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
ndarray	The sliced array along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import slice_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [9, 8, 7, 6, 5]])
>>> out = slice_along_seq(array, start=2)
>>> out
array([[2, 3, 4], [7, 6, 5]])
>>> out = slice_along_seq(array, stop=3)
>>> out
array([[0, 1, 2], [9, 8, 7]])
>>> out = slice_along_seq(array, step=2)
>>> out
array([[0, 2, 4], [9, 7, 5]])

batcharray.array.sort_along_batch

sort_along_batch(
    array: ndarray, kind: SortKind | None = None
) -> ndarray

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

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input 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
ndarray	The sorted array along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import sort_along_batch
>>> array = np.array([[2, 6], [0, 3], [4, 9], [8, 1], [5, 7]])
>>> out = sort_along_batch(array)
>>> out
array([[0, 1], [2, 3], [4, 6], [5, 7], [8, 9]])

batcharray.array.sort_along_seq

sort_along_seq(
    array: ndarray, kind: SortKind | None = None
) -> ndarray

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

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input 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
ndarray	The sorted array along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import sort_along_seq
>>> array = np.array([[7, 3, 0, 8, 5], [1, 9, 6, 4, 2]])
>>> out = sort_along_seq(array)
>>> out
array([[0, 3, 5, 7, 8], [1, 2, 4, 6, 9]])

batcharray.array.split_along_batch

split_along_batch(
    array: ndarray,
    split_size_or_sections: int | Sequence[int],
) -> list[ndarray]

Split the array 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.

Parameters:

Name	Type	Description	Default
array	ndarray	The input 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[ndarray]	The array chunks.

Example usage:

>>> import numpy as np
>>> from batcharray.array import split_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> outputs = split_along_batch(array, split_size_or_sections=2)
>>> outputs
[array([[0, 1], [2, 3]]), array([[4, 5], [6, 7]]), array([[8, 9]])]

batcharray.array.split_along_seq

split_along_seq(
    array: ndarray,
    split_size_or_sections: int | Sequence[int],
) -> list[ndarray]

Split the array 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.

Parameters:

Name	Type	Description	Default
array	ndarray	The input 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[ndarray]	The array chunks.

Example usage:

>>> import numpy as np
>>> from batcharray.array import split_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> outputs = split_along_seq(array, split_size_or_sections=2)
>>> outputs
[array([[0, 1], [5, 6]]), array([[2, 3], [7, 8]]), array([[4], [9]])]

batcharray.array.sum_along_batch

sum_along_batch(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the sum of all elements along the batch axis.

Note

This function assumes the batch axis is the first axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The sum of all elements along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import sum_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = sum_along_batch(array)
>>> out
array([20, 25])
>>> out = sum_along_batch(array, keepdims=True)
>>> out
array([[20, 25]])

batcharray.array.sum_along_seq

sum_along_seq(
    array: ndarray, keepdims: bool = False
) -> ndarray

Return the sum of all elements along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
keepdims	bool	Whether the output array has dim retained or not.	False

Returns:

Type	Description
ndarray	The sum of all elements along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import sum_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = sum_along_seq(array)
>>> out
array([10, 35])
>>> out = sum_along_seq(array, keepdims=True)
>>> out
array([[10], [35]])

batcharray.array.take_along_batch

take_along_batch(
    array: ndarray, indices: ndarray
) -> ndarray

Return a new array which indexes the input array along the batch axis using the entries in indices.

Note

This function assumes the batch axis is the first axis.

Note

Equivalent to index_select_along_batch.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
indices	ndarray	The 1-D array containing the indices to take.	required

Returns:

Type	Description
ndarray	The indexed array along the batch axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import take_along_batch
>>> array = np.array([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]])
>>> out = take_along_batch(array, np.array([2, 4]))
>>> out
array([[4, 5],
       [8, 9]])
>>> out = take_along_batch(array, np.array([4, 3, 2, 1, 0]))
>>> out
array([[8, 9],
       [6, 7],
       [4, 5],
       [2, 3],
       [0, 1]])

batcharray.array.take_along_seq

take_along_seq(array: ndarray, indices: ndarray) -> ndarray

Return a new array which indexes the input array along the sequence axis using the entries in indices.

Note

This function assumes the sequence axis is the second axis.

Note

Equivalent to index_select_along_seq.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
indices	ndarray	The 1-D array containing the indices to take.	required

Returns:

Type	Description
ndarray	The indexed array along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import take_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = take_along_seq(array, np.array([2, 4]))
>>> out
array([[2, 4],
       [7, 9]])
>>> out = take_along_seq(array, np.array([4, 3, 2, 1, 0]))
>>> out
array([[4, 3, 2, 1, 0],
       [9, 8, 7, 6, 5]])

batcharray.array.tile_along_seq

tile_along_seq(array: ndarray, reps: int) -> ndarray

Construct an array by repeating the input array along the sequence axis.

Note

This function assumes the sequence axis is the second axis.

Parameters:

Name	Type	Description	Default
array	ndarray	The input array.	required
reps	int	The number of repetitions data along the sequence axis.	required

Returns:

Type	Description
ndarray	A new array with the data repeated along the sequence axis.

Example usage:

>>> import numpy as np
>>> from batcharray.array import tile_along_seq
>>> array = np.array([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
>>> out = tile_along_seq(array, 2)
>>> out
array([[0, 1, 2, 3, 4, 0, 1, 2, 3, 4],
       [5, 6, 7, 8, 9, 5, 6, 7, 8, 9]])