augmax

library0.4.1pypypi✓ verified 84d ago

Augmax is a library for efficiently composable data augmentation on GPU using JAX. Version 0.4.1 supports JAX transformations and provides a functional, composable augmentation pipeline. Release cadence is occasional; primarily maintained on GitHub.

ai-ml

pip install augmax

github ↗homepage ↗

INSTALL

IMPORT

SIG · AUGMAX

augmax

ai-mlpythonv0.4.1

harness data pending

Install & Compatibility

Where this runs

No compatibility data collected yet for this library.

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

augmax

✓ import augmax

Main library. No common wrong import.

augmax.Chain

✓ from augmax import Chain

✗ import augmax.Chain

Chain is a top-level class; using dot notation incorrectly leads to AttributeError.

augmax.RandomApply

✓ from augmax import RandomApply

Defines a random augmentation pipeline and applies it to a batch of images.

import jax import jax.numpy as jnp import augmax from augmax import Chain, RandomApply, Rotate, FlipX, FlipY # Define a transformation pipeline pipeline = Chain( RandomApply(Rotate(limit=30), p=0.5), RandomApply(FlipX(), p=0.5), RandomApply(FlipY(), p=0.5), ) # Create a random key for JAX key = jax.random.PRNGKey(0) # Create a dummy batch of images (N, H, W, C) images = jnp.ones((4, 256, 256, 3)) # Apply augmentation augmented_images = pipeline(key, images) print(augmented_images.shape) # (4, 256, 256, 3)

Debug

Known issues

gotchaAll augmentations require a JAX PRNG key; using the same key leads to identical augmentations across calls.

fix

Generate a new PRNG key for each call: key, subkey = jax.random.split(key) and pass subkey.

affects: all

breakingaugmax uses the new JAX random key API (PRNGKey). Older JAX versions (<0.4) had different key handling; upgrade JAX if you encounter errors with random keys.

fix

Ensure jax >= 0.4.0 and use jax.random.PRNGKey for key creation.

affects: 0.4+

gotchaThe input batch must be a JAX array (jnp.ndarray); numpy arrays are not directly supported and may cause type errors.

fix

Convert numpy arrays to JAX arrays using jnp.array(...).

affects: all

Errors

Common errors & fixes

AttributeError: module 'augmax' has no attribute 'Chain'

Using import augmax and then augmax.Chain, but Chain is not a submodule.

fix

Use: from augmax import Chain

TypeError: expected a JAX array, got <class 'numpy.ndarray'>

Passing a numpy array instead of a JAX array.

fix

Convert to JAX array: import jax.numpy as jnp; images = jnp.array(numpy_images)

Upgrade

Version history

0.4.1latest on PyPI · released May 16, 2025

Audit

Dependencies

jaxrequiredRequired for GPU/TPU operations and JIT compilation

jaxlibrequiredRequired JAX runtime

Agent activity

22 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/khdlr/augmax ↗

homepagehttps://github.com/khdlr/augmax ↗