parallel_wandb

This simple package makes it easy to use the new (reinit="create_new") feature of Weights & Biases (wandb) to create and log to multiple wandb runs in parallel

This, when combined with jax.vmap, enables extremely efficient, high-throughput training (and logging!) of multiple simultaneous training runs.

This package provides two simple functions that you can import and use in your own project: wandb_init to initialize multiple wandb runs and wandb_log to log metrics to them in parallel.
A demonstration of how these can be used with jax.vmap can be found in jax_mnist.py.

Installation

(optional) Install UV: https://docs.astral.sh/uv/getting-started/installation/
Add this package as a dependency to your project:

uv add parallel_wandb

OR, if you don't use UV yet, you can also pip install parallel_wandb.

Usage

from parallel_wandb import wandb_init, wandb_log

runs = wandb_init(
    {"name": ["run_0", "run_1"], "config": {"seed": [0, 1]}},
    project="test_project",
    name="test_name",
)
assert isinstance(runs, np.ndarray) and runs.shape == (2,) and runs.dtype == object

wandb_log(runs, {"loss": [0.1, 0.2]}, step=0)

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

parallel_wandb-0.0.3.tar.gz (103.3 kB view details)

Uploaded Jan 16, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

parallel_wandb-0.0.3-py3-none-any.whl (17.6 kB view details)

Uploaded Jan 16, 2026 Python 3

File details

Details for the file parallel_wandb-0.0.3.tar.gz.

File metadata

Download URL: parallel_wandb-0.0.3.tar.gz
Upload date: Jan 16, 2026
Size: 103.3 kB
Tags: Source
Uploaded using Trusted Publishing? Yes
Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for parallel_wandb-0.0.3.tar.gz
Algorithm	Hash digest
SHA256	`34578c98eba8ab5d9579bdddc4ac4bec002f6af9e269206b1c17f4732ad907d0`
MD5	`9e772c6910b42123a90e511e34139494`
BLAKE2b-256	`e7cbbec01db9a50e03989d7c5efc9080cb2d58e0cec6462740b5ad4d9b1ff77c`

See more details on using hashes here.

Provenance

The following attestation bundles were made for parallel_wandb-0.0.3.tar.gz:

Publisher: publish.yaml on lebrice/parallel_wandb

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Statement:
- Statement type: https://in-toto.io/Statement/v1
- Predicate type: https://docs.pypi.org/attestations/publish/v1
- Subject name: parallel_wandb-0.0.3.tar.gz
- Subject digest: 34578c98eba8ab5d9579bdddc4ac4bec002f6af9e269206b1c17f4732ad907d0
- Sigstore transparency entry: 832055521
- Sigstore integration time: Jan 16, 2026
Source repository:
- Permalink: lebrice/parallel_wandb@4e9a2fc6e5f5f3f433feb9fe0a9d9f23dc5105ac
- Branch / Tag: refs/tags/v0.0.3
- Owner: https://github.com/lebrice
- Access: public
Publication detail:
- Token Issuer: https://token.actions.githubusercontent.com
- Runner Environment: github-hosted
- Publication workflow: publish.yaml@4e9a2fc6e5f5f3f433feb9fe0a9d9f23dc5105ac
- Trigger Event: release