parajax 0.1.1

Parallelization utilities for JAX

Parajax

Automagic parallelization of calls to JAX-based functions

Features

• 🚀 Device-parallel execution: run across multiple CPUs, GPUs or TPUs automatically
• 🔄 Drop-in replacement for jax.vmap
• ⚡ JIT-compatible: works with jax.jit and variants
• 🪄 Transparent padding when batch sizes aren’t divisible by number of devices
• 🎯 Simple interface: just decorate your function with pvmap

Installation

pip install parajax

Example

import multiprocessing

import jax
import jax.numpy as jnp
from parajax import pvmap

jax.config.update("jax_num_cpu_devices", multiprocessing.cpu_count())
# ^ Only needed on CPU: allow JAX to use all CPU cores

@pvmap
def square(x: float) -> float:
    return x**2

xs = jnp.arange(97)
ys = square(xs)

That's it! Invocations of square will now be automatically parallelized across all available devices.