parajax 0.1.3

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
• ⚡ JIT-compatible: works with jax.jit and variants
• 🪄 Automatic handling of input shapes not divisible by the 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):
    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.

Documentation

For more details, check out the documentation.