Distributed

September 16, 2026 ยท View on GitHub

Device detection and placement, device meshes, SPMD data parallelism and cross-device metric reduction come from substrax, which datarax depends on: substrax.devices, substrax.mesh and substrax.spmd. prefetch_to_device is datarax's own: it lives in datarax.control.prefetcher and is exported from the package root.

Quick Start

import jax
from substrax.devices import get_batch_size_recommendation
from substrax.mesh import DeviceMeshManager

# Check available devices
print(f"Devices: {jax.devices()}")

# Get a batch-size recommendation for the detected hardware
recommendation = get_batch_size_recommendation()
print(f"Recommended batch size: {recommendation.optimal_batch_size}")

# Create a 2D mesh for data + model parallelism
mesh = DeviceMeshManager.create_device_mesh({"data": 2, "model": 4})

Multi-Host Training

For multi-host setups, JaxProcessSharderModule derives the shard topology from Grain's ShardByJaxProcess, so each process automatically slices its local shard:

# Each host runs this code
from datarax.sharding import JaxProcessSharderModule

sharder = JaxProcessSharderModule()

# Shard data across hosts (process index/count are auto-derived)
local_batch = sharder.shard_data(global_batch)

See Also