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
Metadata
1 maintainer
Release files for parajax 0.3.4
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| parajax-0.3.4.tar.gz | 4.3 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| parajax-0.3.4-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 9.1 kB
Release files / parajax-0.3.4.tar.gz
| Download URL | parajax-0.3.4.tar.gz |
|---|---|
| Size | 4.3 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
196f85a86a14d76d2e3443556d5c790a1b5757601ed5684aae0cb1a3f01937f5
|
|
BLAKE2b-256 checksum How to use checksums |
3a741e4f8fd58db531bb7864de509794289aca1e4c5b365739d1da7b1ffcf828
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
uv/0.12.7 {"installer":{"name":"uv","version":"0.12.7","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}
|
Release files / parajax-0.3.4-py3-none-any.whl
| Download URL | parajax-0.3.4-py3-none-any.whl |
|---|---|
| Size | 4.8 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
5e9864b1adca480ab62d1ecf5c31f667ebc520c39cdf021c37969728d025f022
|
|
BLAKE2b-256 checksum How to use checksums |
027f18c6630de471c690fd580b890ac4dbbfdb058a7ac861b1e06f2bb98b51bc
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
uv/0.12.7 {"installer":{"name":"uv","version":"0.12.7","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}
|
