axis-batch 0.1.1

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

Axis Batch

---

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

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