Devices

substrax.devices answers two questions: what hardware is this process running on, and where should a batch go.

from substrax.devices import DevicePlacement, detect_devices, get_batch_size_recommendation

info = detect_devices()          # platform, kind, count, one device_kind per device
placement = DevicePlacement()    # defaults to jax.devices()
batch = placement.place_on_device(batch)
recommendation = get_batch_size_recommendation()  # for the detected hardware

DeviceInfo is a frozen dataclass; detect_devices() reads it from jax.devices() and never imports an accelerator plugin itself. DevicePlacement carries the batch-size recommendation table per HardwareType.

jaxlib ships no type stubs for jax.Device, so it cannot annotate a parameter. DeviceLike is the protocol substrax annotates devices with (id, platform, device_kind), which a jax.Device satisfies, and visible_devices() returns jax.devices() typed as it.

substrax.devices

Device information and device placement.

DeviceInfo(*, platform, kind, count, device_kinds) dataclass

Snapshot of the devices visible to the current JAX process.

Attributes:

Name	Type	Description
platform	str	The default backend name exactly as JAX reports it.
kind	DeviceKind	The accelerator class of that backend.
count	int	The number of visible devices.
device_kinds	tuple[str, ...]	One device_kind string per visible device, in JAX order.

has_accelerator property

Whether the default backend is anything other than the CPU.

DeviceKind

Bases: StrEnum

The accelerator class behind a JAX backend name.

from_platform(platform) classmethod

Map a backend name as jax.default_backend() reports it onto a kind.

Parameters:

Name	Type	Description	Default
platform	str	The backend name, in any case ("gpu", "cuda", "METAL").	required

Returns:

Type	Description
DeviceKind	The matching kind, or OTHER for a backend this package has not seen.

DeviceLike

Bases: Protocol

The part of a jax.Device substrax reads.

jaxlib ships no type stubs for Device (jax.Device is a runtime alias of an extension type), so jax.Device cannot annotate a parameter; this protocol types one instead, and a jax.Device satisfies it.

id property

The device's index among its backend's devices.

platform property

The backend the device belongs to, such as "cpu" or "gpu".

device_kind property

The runtime's name for the hardware, such as "NVIDIA H200".

BatchSizeRecommendation(min_batch_size, optimal_batch_size, critical_batch_size, max_memory_batch_size=None, notes='') dataclass

Hardware-specific batch size recommendations.

Attributes:

Name	Type	Description
min_batch_size	int	Minimum batch size for reasonable efficiency.
optimal_batch_size	int	Optimal batch size for peak throughput.
critical_batch_size	int	Critical batch size for reaching roofline (per JAX guide).
max_memory_batch_size	int | None	Maximum batch size before OOM (estimate).
notes	str	Additional notes about the recommendation.

DevicePlacement(default_device=None)

Utility class for explicit device placement of JAX arrays.

This class provides methods for placing arrays on specific devices, distributing batches across devices using sharding, and providing hardware-aware batch size recommendations.

Example

placement = DevicePlacement()
data = jnp.ones((4, 8))
gpu_data = placement.place_on_device(data, jax.devices("gpu")[0])  # Place on first GPU
rec = placement.get_batch_size_recommendation()  # Get batch size recommendation
print(f"Optimal batch: {rec.optimal_batch_size}")

Initialize DevicePlacement.

Parameters:

Name	Type	Description	Default
default_device	DeviceLike | None	Default device to use when none is specified. If None, lazily resolves to jax.devices()[0] on first access.	None

default_device property

Get the default device, lazily resolving on first access.

hardware_type property

Get the detected hardware type, lazily resolving on first access.

num_devices property

Get the number of available devices.

place_on_device(data, device=None)

Place data on a specific device.

This uses jax.device_put for explicit device placement, ensuring data is transferred to the target device.

Parameters:

Name	Type	Description	Default
data	PyTree	PyTree of JAX arrays to place on device.	required
device	DeviceLike | None	Target device. If None, uses the default device.	None

Returns:

Type	Description
PyTree	PyTree with arrays placed on the specified device.

Example

data = {"images": jnp.ones((4, 28, 28, 3))}
gpu_data = placement.place_on_device(data, jax.devices("gpu")[0])

replicate_across_devices(data, devices=None)

Replicate data across all specified devices.

Creates a copy of the data on each device, useful for broadcasting model weights or constants.

Parameters:

Name	Type	Description	Default
data	PyTree	PyTree of JAX arrays to replicate.	required
devices	Sequence[DeviceLike] | None	List of devices to replicate to. If None, uses all devices.	None

Returns:

Type	Description
PyTree	PyTree with arrays replicated across devices.

shard_batch_dim(data, mesh, batch_axis=0, mesh_axis='data')

Shard data along the batch dimension.

This is the most common sharding pattern for data-parallel training, where each device processes a slice of the batch.

Parameters:

Name	Type	Description	Default
data	PyTree	PyTree of JAX arrays to shard.	required
mesh	Mesh	Device mesh to shard across.	required
batch_axis	int	The axis index representing the batch dimension.	0
mesh_axis	str	The mesh axis name to shard along.	'data'

Returns:

Type	Description
PyTree	PyTree with arrays sharded along the batch dimension.

get_batch_size_recommendation(hardware_type=None)

Get batch size recommendation for the current hardware.

Parameters:

Name	Type	Description	Default
hardware_type	HardwareType | None	Override hardware type. If None, uses detected type.	None

Returns:

Type	Description
BatchSizeRecommendation	BatchSizeRecommendation with hardware-specific values.

validate_batch_size(batch_size, warn_suboptimal=True)

Validate batch size against hardware recommendations.

Parameters:

Name	Type	Description	Default
batch_size	int	The batch size to validate.	required
warn_suboptimal	bool	Whether to warn for suboptimal (but valid) sizes.	True

Returns:

Type	Description
tuple[bool, str]	Tuple of (is_valid, message).

HardwareType

Bases: Enum

Enumeration of supported hardware types.

detect_devices()

Probe the JAX runtime once and return what it reports.

Errors from the runtime (no backend, uninitialised plugins) propagate; a process that cannot see its devices should not be told it is running on a CPU.

Returns:

Type	Description
DeviceInfo	The device snapshot.

visible_devices()

The devices jax.devices() reports, typed as :class:DeviceLike.

Returns:

Type	Description
Sequence[DeviceLike]	Every device of the default backend, in JAX's order.

get_batch_size_recommendation(hardware_type=None)

Get batch size recommendation for current or specified hardware.

Parameters:

Name	Type	Description	Default
hardware_type	HardwareType | None	Hardware type to get recommendation for.	None

Returns:

Type	Description
BatchSizeRecommendation	BatchSizeRecommendation with hardware-specific values.

place_on_device(data, device=None)

Convenience function for placing data on a device.

Parameters:

Name	Type	Description	Default
data	PyTree	PyTree of JAX arrays.	required
device	DeviceLike | None	Target device. If None, uses first available device.	None

Returns:

Type	Description
PyTree	PyTree with arrays on the specified device.