augmax

What it is and what it does

Augmax is a JAX-native image augmentation library designed to apply sequences of random transformations to images with minimal memory overhead. Unlike traditional augmentation frameworks that execute each transformation separately, Augmax fuses operations together to reduce redundant memory reads and writes. It integrates with JAX's jit compilation and vmap vectorization, enabling efficient batch processing on GPUs and TPUs.

The package is intended for machine learning workflows where data augmentation is a bottleneck. You define a pipeline as a chain of transformations (similar to Albumentations), then apply it to single images or entire batches. JAX's functional programming model allows Augmax to compile the entire pipeline into a single optimized kernel, making it particularly useful for large-scale training on accelerators.

Use it for:

• Accelerate training pipelines by fusing image augmentations into a single GPU kernel via jax.jit.
• Apply different random augmentations to each image in a batch using jax.vmap without loop overhead.
• Build reproducible augmentation pipelines with JAX's deterministic random number generation.
• Reduce memory bandwidth during data loading by composing transformations before GPU transfer.
• Integrate augmentation directly into JAX training loops without leaving the JAX ecosystem.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Augmax is a JAX-based image data augmentation framework that chains transformations together and compiles them efficiently for GPU/TPU execution, minimizing redundant memory operations.

Yes, if you are training models with JAX and need efficient GPU-accelerated augmentation. The low install friction, permissive license, and active repository make it a solid choice. However, maintenance is aging (last update 455 days ago), so verify that the transformation set and API meet your needs before committing to it in production. No known vulnerabilities.

Install

augmax on PyPI

pip install augmax

uv add augmax

poetry add augmax

Installing augmax

Before you install

Low friction install with a pure Python wheel. Maintenance is aging—last commit was 2025-05-16 and the project has not been updated for 455 days—but the repository remains active and not archived. Depends on einops, jax, and numpy.

License in practice

Licensed under Apache-2.0 (permissive), which allows commercial and private use with minimal restrictions. No licensing friction for most use cases.

Quickstart

pip install augmax

import jax
import augmax

transform = augmax.Chain(
    augmax.RandomCrop(256, 256),
    augmax.HorizontalFlip(),
)
rng = jax.random.PRNGKey(27)
transformed_image = jax.jit(transform)(rng, image)

Requires JAX and its dependencies (including a compatible CUDA/GPU setup if GPU acceleration is desired); numpy and einops must be installed.

Verify before relying

• Exact list of supported augmentation operations beyond RandomCrop, HorizontalFlip, and Rotate.
• Performance benchmarks comparing Augmax to other frameworks (e.g., Albumentations) on typical workloads.
• Whether the API is fully compatible with Albumentations as claimed, or only partially similar.

Package facts

Evidence: augmax-0.4.1-py3-none-any.whl

Tags

More Artificial Intelligence packages