Welcome to Axis Batch documentation!

Installation & Usage

Project Info

API Reference

Axis Batch

CI Status Documentation Status Test coverage percentage

uv Ruff pre-commit

PyPI Version Supported Python versions License

Documentation: https://axis-batch.readthedocs.io

Source Code: https://github.com/34j/axis-batch

Divide an array into batches along specific axes in NumPy / PyTorch / JAX

Installation

Install this via pip (or your favourite package manager):

pip install axis-batch

Usage

import numpy as np

from axis_batch import AxisBatch

a = np.arange(12).reshape(3, 4)
b = AxisBatch(a, axis=0, size=2)
for i, x in enumerate(b):
    print(f"{i}: {x}")
    b.send(x + 1)
print(b.value)

0: [[0 1 2 3]
 [4 5 6 7]]
1: [[ 8  9 10 11]]
[[ 1  2  3  4]
 [ 5  6  7  8]
 [ 9 10 11 12]]

Contributors ✨

Thanks goes to these wonderful people (emoji key):

This project follows the all-contributors specification. Contributions of any kind welcome!

Credits

Copier

This package was created with Copier and the browniebroke/pypackage-template project template.