parajax 0.2.2

Parallelization utilities for JAX

Automagic parallelization of calls to JAX-based functions

Features

• 🚀 Device-parallel execution: run across multiple CPUs, GPUs or TPUs automatically
• ⚡ Fully composable with jax.jit, jax.vmap, and other JAX transformations
• 🪄 Automatic handling of input shapes not divisible by the number of devices
• 🎯 Simple interface: just decorate your function with autopmap

Installation

pip install parajax

Example

import multiprocessing

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

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

@autopmap
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.