GitHub repository: 34j/axis-batch
Primary language: Python
Description: Divide an array into batches along specific axes in NumPy / PyTorch / JAX