batcharray

Overview

batcharray is lightweight library built on top of NumPy to manipulate nested data structure with NumPy arrays. This library provides functions for arrays where the first axis is the batch axis. It also provides functions for arrays representing a batch of sequences where the first axis is the batch axis and the second axis is the sequence axis.

• Motivation
• Documentation
• Installation
• Features
• Contributing
• API stability
• License

Motivation

Let's imagine you have a batch which is represented by a dictionary with three arrays, and you want to take the first 2 items. batcharray provides the function slice_along_batch that allows to slide all the arrays:

>>> import numpy as np
>>> from batcharray.nested import slice_along_batch
>>> batch = {
...     "a": np.array([[2, 6], [0, 3], [4, 9], [8, 1], [5, 7]]),
...     "b": np.array([4, 3, 2, 1, 0]),
...     "c": np.array([1.0, 2.0, 3.0, 4.0, 5.0]),
... }
>>> slice_along_batch(batch, stop=2)
{'a': array([[2, 6], [0, 3]]), 'b': array([4, 3]), 'c': array([1., 2.])}

Similarly, it is possible to split a batch in multiple batches by using the function split_along_batch:

>>> import numpy as np
>>> from batcharray.nested import split_along_batch
>>> batch = {
...     "a": np.array([[2, 6], [0, 3], [4, 9], [8, 1], [5, 7]]),
...     "b": np.array([4, 3, 2, 1, 0]),
...     "c": np.array([1.0, 2.0, 3.0, 4.0, 5.0]),
... }
>>> split_along_batch(batch, split_size_or_sections=2)
[{'a': array([[2, 6], [0, 3]]), 'b': array([4, 3]), 'c': array([1., 2.])},
 {'a': array([[4, 9], [8, 1]]), 'b': array([2, 1]), 'c': array([3., 4.])},
 {'a': array([[5, 7]]), 'b': array([0]), 'c': array([5.])}]

Please check the documentation to see all the implemented functions.

Features

batcharray provides several modules for different use cases:

• array: Functions to manipulate single NumPy arrays with batch and sequence dimensions
  • Slicing, indexing, and selection operations
  • Statistical reductions (mean, max, min, etc.)
  • Sorting and permutation operations
  • Cumulative operations (cumsum, cumprod)
  • Support for both regular and masked arrays
• nested: Functions to manipulate nested data structures (dicts, lists) containing arrays
  • Apply operations to all arrays in a structure simultaneously
  • Maintain relationships between related arrays
  • Mathematical operations (abs, exp, log, trigonometric functions)
  • All operations from array module, but for nested structures
• computation: Computation models for abstracting array operations
  • Automatic selection of appropriate implementation based on array type
  • Support for regular arrays and masked arrays
  • Extensible for custom array types
• recursive: Tools for recursively applying functions to nested structures
  • Traverse complex hierarchies of dicts, lists, and tuples
  • Apply transformations to specific types
  • Customizable traversal strategies
• utils: Utility functions for exploring nested structures
  • Breadth-first search (BFS) for wide structures
  • Depth-first search (DFS) for deep structures
  • Helpful for debugging and validation

Documentation

• latest (stable): documentation from the latest stable release.
• main (unstable): documentation associated to the main branch of the repo. This documentation may contain a lot of work-in-progress/outdated/missing parts.

Installation

We highly recommend installing a virtual environment. batcharray can be installed from pip using the following command:

uv pip install batcharray

To make the package as slim as possible, only the minimal packages required to use batcharray are installed. To include all the dependencies, you can use the following command:

uv pip install batcharray[all]

Please check the get started page to see how to install only some specific dependencies or other alternatives to install the library. The following is the corresponding batcharray versions and tested dependencies.

^* indicates an optional dependency

Contributing

Please check the instructions in CONTRIBUTING.md.

Suggestions and Communication

Everyone is welcome to contribute to the community. If you have any questions or suggestions, you can submit Github Issues. We will reply to you as soon as possible. Thank you very much.

API stability

:warning: While batcharray is in development stage, no API is guaranteed to be stable from one release to the next. In fact, it is very likely that the API will change multiple times before a stable 1.0.0 release. In practice, this means that upgrading batcharray to a new version will possibly break any code that was using the old version of batcharray.

License

batcharray is licensed under BSD 3-Clause "New" or "Revised" license available in LICENSE file.