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
@parallelize
Installation
pip install parajax
Example
import multiprocessing
import jax
import jax.numpy as jnp
from parajax import parallelize
jax.config.update("jax_num_cpu_devices", multiprocessing.cpu_count())
# ^ Only needed on CPU: allow JAX to use all CPU cores
@parallelize
def square(xs):
return xs**2
xs = jnp.arange(12_345)
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.
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
parajax-0.3.3.tar.gz
(4.3 kB
view details)
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file parajax-0.3.3.tar.gz.
File metadata
- Download URL: parajax-0.3.3.tar.gz
- Upload date:
- Size: 4.3 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
uv/0.12.1 {"installer":{"name":"uv","version":"0.12.1","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
8c44652d5bc6537b9291d45de1fd1181e6ae411ec92ed0a83de2455648b3318d
|
|
| MD5 |
10d3a4b878729de0c4e1339c94acd753
|
|
| BLAKE2b-256 |
8d602303f8045feac76229c9d744fc51944ae1876753ae7604d77f53c278c9db
|
File details
Details for the file parajax-0.3.3-py3-none-any.whl.
File metadata
- Download URL: parajax-0.3.3-py3-none-any.whl
- Upload date:
- Size: 4.8 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
uv/0.12.1 {"installer":{"name":"uv","version":"0.12.1","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
627025dee12652e832dca44deefcb8595afff9571e349834e6b167825c0167b0
|
|
| MD5 |
b004ed2e7f58e42bd064d4aeded9d4ed
|
|
| BLAKE2b-256 |
0c8f20b8e0337ee22a3f8d52a0185519092ac138028da6c00f343ab0cca5e39b
|
